fix differentiation of loopinfo exprs

addresses the first part of FluxML#897
simeonschaub · Apr 26, 2021 · 3218e50 · 3218e50
1 parent 4e54eb7
commit 3218e50
Show file tree

Hide file tree

Showing 3 changed files with 5 additions and 2 deletions.
diff --git a/src/compiler/reverse.jl b/src/compiler/reverse.jl
@@ -275,7 +275,7 @@ function adjoint(pr::Primal)
         end
       elseif ex isa Core.PiNode
         grads[ex.val] = grads[v]
-      elseif isexpr(ex, GlobalRef, :call, :isdefined, :inbounds, :meta)
+      elseif isexpr(ex, GlobalRef, :call, :isdefined, :inbounds, :meta, :loopinfo)
       elseif isexpr(ex)
         push!(rb, stmt(xcall(Base, :error, "Can't differentiate $(ex.head) expression"),
                        line = b[v].line))

diff --git a/test/compiler.jl b/test/compiler.jl
@@ -143,3 +143,6 @@ end
   ms = MyStruct(1, 2)
   @test Zygote.gradient(sumall, ms) == ((a = 2, b = 2),)
 end
+
+# issue 897
+@test gradient(x -> sum(norm, collect(eachcol(x))), ones(3, 400))[1] ≈ fill(0.5773502691896258, 3, 400)
diff --git a/test/features.jl b/test/features.jl
@@ -402,7 +402,7 @@ function pow_simd(x, n)
   return r
 end
 
-@test_broken gradient(pow_simd, 2, 3) == (12,nothing)
+@test gradient(pow_simd, 2, 3) == (12,nothing)
 
 @testset "tuple getindex" begin
   @test gradient(x -> size(x)[2], ones(2,2,2)) == (nothing,)