Merge #1079

1079: ProjectTo{AbstractArray}(::Tangent{Any}) fix r=willtebbutt a=willtebbutt Resolves #1078 Co-authored-by: WT <[email protected]> Co-authored-by: willtebbutt <[email protected]>
FluxML · Sep 24, 2021 · 446594d · 446594d · willtebbutt · Sep 24, 2021
2 parents 700c314 + 79454f3
commit 446594d
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "Zygote"
 uuid = "e88e6eb3-aa80-5325-afca-941959d7151f"
-version = "0.6.22"
+version = "0.6.23"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

diff --git a/src/compiler/chainrules.jl b/src/compiler/chainrules.jl
@@ -150,6 +150,12 @@ _project(x::AbstractArray, dx::Tuple) = _project(x, reshape(collect(dx), axes(x)
 # CRC likes Tangent{<:Complex}, but Zygote makes Tangent{Any}
 (project::ProjectTo{<:Complex})(dx::Tangent) = project(Complex(dx.re, dx.im))
 
+# CRC likes Tangent{AbstractArray}, but Zygote makes Tangent{Any}
+# in particular this would hit https://github.com/JuliaDiff/ChainRulesCore.jl/blob/2ec2549b73b22bc08f554dae864fb650cfb9c3d7/src/projection.jl#L139
+# if we were not losing track of the Primal in the Tangent
+# This type piracy is just giving up that safety check.
+(project::ProjectTo{AbstractArray})(dx::Tangent) = dx
+
 """
   ZBack{F}(back) <: Function
 

diff --git a/test/chainrules.jl b/test/chainrules.jl
@@ -335,6 +335,13 @@ end
         test_rrule(ZygoteRuleConfig(), sum, x -> cbrt(x), randn(5))
         test_rrule(ZygoteRuleConfig(), sum, x -> cbrt(x), randn(5); rrule_f=rrule_via_ad)
     end
+
+    # See https://github.com/FluxML/Zygote.jl/issues/1078
+    @testset "ProjectTo{AbstractArray}(::Tangent{Any})" begin
+        X = UpperHessenberg(randn(5, 5))
+        dX = Tangent{Any}(element=randn(5, 5))
+        @test ProjectTo(X)(dX) === dX
+    end
 end
 
 @testset "FastMath support" begin