ERROR: Can't differentiate simdloop expression

[baggepinnen]

The following example (from ForwardDiff.jl readme) works well on small input vectors but fails on larger ones.

using Zygote
f(x) = sum(sin, x) + prod(tan, x) * sum(sqrt, x);
gz = x -> Zygote.gradient(f, x)[1]; # g = ∇f

x = rand(5); # small size works well
gz(x)
x = rand(5000); # large size errors
julia> gz(x)
ERROR: Can't differentiate simdloop expression
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] macro expansion at ./simdloop.jl:71 [inlined]
 [3] mapreduce_impl at ./reduce.jl:164 [inlined]
 [4] (::typeof(∂(Base.mapreduce_impl)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [5] mapreduce_impl at ./reduce.jl:173 [inlined]
 [6] (::typeof(∂(Base.mapreduce_impl)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [7] mapreduce_impl at ./reduce.jl:173 [inlined]
 [8] (::typeof(∂(Base.mapreduce_impl)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [9] mapreduce_impl at ./reduce.jl:173 [inlined]
 [10] (::typeof(∂(Base.mapreduce_impl)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [11] mapreduce_impl at ./reduce.jl:178 [inlined]
 [12] (::typeof(∂(Base.mapreduce_impl)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [13] _mapreduce at ./reduce.jl:320 [inlined]
 [14] (::typeof(∂(Base._mapreduce)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [15] _mapreduce_dim at ./reducedim.jl:308 [inlined]
 [16] (::typeof(∂(Base._mapreduce_dim)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [17] #mapreduce#548 at ./reducedim.jl:304 [inlined]
 [18] mapreduce at ./reducedim.jl:304 [inlined]
 [19] _prod at ./reducedim.jl:653 [inlined]
 [20] #prod#553 at ./reducedim.jl:649 [inlined]
 [21] prod at ./reducedim.jl:649 [inlined]
 [22] f at ./REPL[3]:1 [inlined]
 [23] (::typeof(∂(f)))(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface2.jl:0
 [24] (::getfield(Zygote, Symbol("##71#72")){typeof(∂(f))})(::Float64) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface.jl:38
 [25] gradient(::Function, ::Array{Float64,1}) at /local/home/fredrikb/.julia/packages/Zygote/7kv2i/src/compiler/interface.jl:47
 [26] (::getfield(Main, Symbol("##3#4")))(::Array{Float64,1}) at ./REPL[4]:1
 [27] top-level scope at none:0

If, however, the definition of f is changed to
f(x) = sum(sin.(x)) + prod(tan.(x)) * sum(sqrt.(x))
it works

[MikeInnes]

We should have adjoints for sum(f, x) and prod(f, x) (even if they just do that rewrite internally). Which of those is the issue, or is it both?

Separately we can probably handle simdloop expressions, even just by ignoring it, so we should fix that.

[baggepinnen]

prod(tan, x) is the problem, sum works

[MikeInnes]

Alright, fixed that specific issue. Will leave this open for the simdloop component.

[dillondaudert]

Not sure if this is fixed on master (couldn't get it to precompile on Julia v1.1), but I get this error on Zygote 0.2 when calling gradient on evaluate from Distances.jl.

using Zygote
using Distances
gradient((a, b) -> evaluate(Euclidean(), a, b), rand(10), rand(10))

[MikeInnes]

Yup, coming from this line. Given that evaluate is basically mapreduce you could probably get away with a custom adjoint here pretty easily.

[GiggleLiu]

Just one more example to show this kind of problem can be quite dangerous.
Somtimes simdloop happens unexpectedly, even if your program is tested. Waiting for the simdloop component.

julia> using Zygote
julia> using Statistics: var
julia> using LinearAlgebra: norm

julia> function loss1(x)
           l1_pairs = [(1,5), (2,6), (3,7), (4,8)]
           var([norm(x[i]-x[j]) for (i, j) in l1_pairs])
       end
loss1 (generic function with 1 method)

julia> loss1'(randn(2,8))
2×8 Array{Float64,2}:
 -0.548216  -0.76573   0.548216  0.76573   0.0  0.0  0.0  0.0
 -0.801852  -0.584338  0.801852  0.584338  0.0  0.0  0.0  0.0

julia> function loss2(x)  # similar function, the length of pairs are different.
           l2_pairs = [(1,2), (2,3), (3,4), (4,1), (5,6), (6,7), (7,8), (8,5),
                       (1,6), (1, 8), (2,7), (2,5), (3,6), (3,8), (4,7),(4,5)]
           var([norm(x[i]-x[j]) for (i, j) in l2_pairs])
       end
loss2 (generic function with 1 method)

julia> loss2'(randn(2,8))
ERROR: Can't differentiate simdloop expression
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] macro expansion at ./simdloop.jl:75 [inlined]
....

[MikeInnes]

In most cases like this we should just add gradients for var, norm etc, or whatever's causing this.

Of course it'd be nice if it worked out of the box. I don't think simd loops are hard to differentiate -- we probably just ignore the loopinfo/simdloop metadata entirely, in which case it's trivial -- it just needs someone to look over it and verify that.

[baggepinnen]

I tried all pasted code in this issue and no problems persist, except for the code with Distances which now suffers from #385

[haozhangphd]

The same error message resurfaced with the above code on Julia 1.6, but not on Julia 1.5. See #897.