Merge pull request #1102 from mcabbott/ref_tangent

Make `wrap_chainrules_input` understand `mutable struct`s

Project.toml

@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.27"
+version = "0.6.28"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
@@ -24,7 +24,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 AbstractFFTs = "0.5, 1.0"
 ChainRules = "1.5"
-ChainRulesCore = "1.6"
+ChainRulesCore = "1.9"
 ChainRulesTestUtils = "1"
 DiffRules = "1.0"
 FillArrays = "0.8, 0.9, 0.10, 0.11, 0.12"

src/compiler/chainrules.jl

@@ -128,6 +128,8 @@ Convert `x` from the format Zygote uses internally to differentials types ChainR
   xp = map(wrap_chainrules_input, xs)
   ChainRules.Tangent{Any, typeof(xp)}(xp)
 end
+# For mutable types, including x=Ref(1), Zygote makes Ref{Any}(::NamedTuple)
+@inline wrap_chainrules_input(x::Ref) = wrap_chainrules_input(x[])
 
 """
   _project(x, dx)

test/features.jl

@@ -443,6 +443,20 @@ let
   @test back(1.) == ((1.0,),)
 end
 
+@testset "mutable struct, including Ref" begin
+  # Zygote's representation is Base.RefValue{Any}((value = 7.0,)), but the
+  # map to ChainRules types and back normalises to (value = 7.0,) same as struct:
+  @test gradient(x -> x.value^2 + x.value, MyMutable(3)) === ((value = 7.0,),)
+
+  # Same for Ref. This doesn't seem to affect `pow_mut` test in this file.
+  @test gradient(x -> x.x^2 + x.x, Ref(3)) === ((x = 7.0,),)
+  @test gradient(x -> real(x.x^2 + im * x.x), Ref(4)) === ((x = 8.0,),)
+
+  # Broadcasting over Ref is handled specially. Tested elsehwere too.
+  @test gradient(x -> sum(sum, x .* [1,2,3]), Ref([4,5])) == ((x = [6.0, 6.0],),)
+  @test gradient(x -> sum(sum, Ref(x) .* [1,2,3]), [4,5]) == ([6.0, 6.0],)
+end
+
 function type_test()
    Complex{<:Real}
 end