Open doddgray opened 3 years ago
Shorter way to get the same error:
julia> gradient(x -> sum(convert(SMatrix{2,2,Float64,4}, x)), rand(2,2))
ERROR: Need an adjoint for constructor SMatrix{2, 2, Float64, 4}. Gradient is
of type FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}
julia> Zygote.@adjoint convert(T::Type{<:SMatrix}, x::AbstractMatrix) = T(x), dv -> (nothing, dv)
julia> gradient(x -> sum(convert(SMatrix{2,2,Float64,4}, x)), rand(2,2))
(2×2 Fill{Float64}: entries equal to 1.0,)
julia> Zygote.gradient(x->sum(f3(x)),3.3)
(6.887761171372725,)
Not entirely sure that defining an adjoint for convert
can't be avoided, though.
thanks for taking a look at this and giving the example convert
adjoint rule. Based on the different lines of StaticArrays/src/convert.jl
referenced in the stack traces from the errors handled and not-handled by the constructor rules, it seems related to static sizing?
Using the same f2(x)
and f3(x)
defs as above:
julia> f2(x) = sum(SMatrix{2,2,Float64,4}(x,x^2,√x,sin(x)))
julia> gradient(f2,3.3) # before SMatrix constructor rule definition
ERROR: Need an adjoint for constructor SMatrix{2, 2, Float64, 4}. Gradient is of type FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:33
[2] (::Zygote.Jnew{SMatrix{2, 2, Float64, 4}, Nothing, false})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/lib/lib.jl:314
[3] (::Zygote.var"#1729#back#202"{Zygote.Jnew{SMatrix{2, 2, Float64, 4}, Nothing, false}})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
[4] Pullback
@ ~/.julia/packages/StaticArrays/w5a7P/src/SArray.jl:23 [inlined]
[5] (::typeof(∂(SMatrix{2, 2, Float64, 4})))(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface2.jl:0
[6] Pullback
@ ~/.julia/packages/StaticArrays/w5a7P/src/convert.jl:4 [inlined]
[7] Pullback
@ ./REPL[6]:1 [inlined]
[8] (::typeof(∂(f2)))(Δ::Float64)
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface2.jl:0
[9] (::Zygote.var"#43#44"{typeof(∂(f2))})(Δ::Float64)
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface.jl:41
[10] gradient(f::Function, args::Float64)
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface.jl:59
[11] top-level scope
@ REPL[10]:1
julia> ChainRulesCore.rrule(T::Type{<:SMatrix}, x::AbstractMatrix) = ( T(x), dv -> (NO_FIELDS, dv) )
julia> ChainRulesCore.rrule(T::Type{<:SMatrix}, xs::Number...) = ( T(xs...), dv -> (NO_FIELDS, dv...) )
julia> Zygote.refresh()
julia> gradient(f2,3.3)
(6.887761171372725,)
julia> function f3(x)::SMatrix{2,2,Float64,4}
[ x x^2; √x sin(x) ]
end
julia> Zygote.gradient(x->sum(f3(x)),3.3)
ERROR: Need an adjoint for constructor SMatrix{2, 2, Float64, 4}. Gradient is of type FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}}
Stacktrace:
[1] error(s::String)
@ Base ./error.jl:33
[2] (::Zygote.Jnew{SMatrix{2, 2, Float64, 4}, Nothing, false})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/lib/lib.jl:314
[3] (::Zygote.var"#1729#back#202"{Zygote.Jnew{SMatrix{2, 2, Float64, 4}, Nothing, false}})(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ~/.julia/packages/ZygoteRules/OjfTt/src/adjoint.jl:59
[4] Pullback
@ ~/.julia/packages/StaticArrays/w5a7P/src/SArray.jl:23 [inlined]
[5] (::typeof(∂(SMatrix{2, 2, Float64, 4})))(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ./compiler/interface2.jl:0
[6] Pullback
@ ~/.julia/packages/StaticArrays/w5a7P/src/convert.jl:36 [inlined]
[7] (::typeof(∂(_convert)))(Δ::FillArrays.Fill{Float64, 2, Tuple{Base.OneTo{Int64}, Base.OneTo{Int64}}})
@ Zygote ./compiler/interface2.jl:0
[8] Pullback
@ ~/.julia/packages/StaticArrays/w5a7P/src/convert.jl:33 [inlined]
[9] Pullback
@ ./REPL[13]:1 [inlined]
[10] (::typeof(∂(#7)))(Δ::Float64)
@ Zygote ./compiler/interface2.jl:0
[11] (::Zygote.var"#43#44"{typeof(∂(#7))})(Δ::Float64)
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface.jl:41
[12] gradient(f::Function, args::Matrix{Float64})
@ Zygote ~/.julia/packages/Zygote/6HN9x/src/compiler/interface.jl:59
[13] top-level scope
@ REPL[13]:1
julia> gradient(x -> sum(convert(SMatrix{2,2,Float64,4}, x)), rand(2,2))
# ... same error as with f3(x) above
The former case (handled by rrule
) references
whereas the latter case references
Maybe this would work if the unroll_tuple(a,l)
was splatted in the SA
constructor? I'll have to check
Why does this fail
while the following alternatives work?
I see the same behavior with
SMatrix
constructor adjoints defined usingChainRulesCore.rrule
or usingZygote.@adjoint
as, respectivelyor
It seems like function defs with return-type specifications cause Zygote missing-constructor-adjoint errors even when the constructor adjoint/rrule is defined, but I could be misunderstanding something more basic.
Per a helpful suggestion from @oxinabox on the Julia/#autodiff Slack channel, I also tried writing a rule with all of the type parameter specified:
but this didn't seem to have any effect. Is there some reason
convert
would be missing custom rules? What's a programmer to do?