Merge pull request #109 from glennmoy/gm/chainrules-bugfic

Fix ChainRules Ambiguity Error
mcabbott · May 25, 2022 · 94dcb6d · 94dcb6d · mcabbott · May 25, 2022
2 parents f172d7a + 3e7bcf5
commit 94dcb6d
Show file tree

Hide file tree

Showing 4 changed files with 7 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "AxisKeys"
 uuid = "94b1ba4f-4ee9-5380-92f1-94cde586c3c5"
 license = "MIT"
-version = "0.2.2"
+version = "0.2.3"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"

diff --git a/src/chainrules.jl b/src/chainrules.jl
@@ -9,6 +9,7 @@ function ChainRulesCore.ProjectTo(x::KeyedArray)
 end
 
 (project::ProjectTo{KeyedArray})(dx) = wrapdims(project.data(parent(dx)), project.keys...)
+(project::ProjectTo{KeyedArray})(dx::AbstractZero) = dx
 
 _KeyedArray_pullback(ȳ, project) = (NoTangent(), project(ȳ))
 _KeyedArray_pullback(ȳ::Tangent, project) = _KeyedArray_pullback(ȳ.data, project)

diff --git a/test/_chainrules.jl b/test/_chainrules.jl
@@ -18,10 +18,14 @@
         p = ProjectTo(ka)
         @test p(data) == ka
 
+        @test p(NoTangent()) == NoTangent()
+
         data = rand(3, 4)
         ka = wrapdims(data, a=1:3, b='a':'d')
         p = ProjectTo(ka)
         @test p(data) == ka
+
+        @test p(NoTangent()) == NoTangent()
     end
 
     @testset "KeyedVector" begin

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -1,6 +1,6 @@
 using Test, AxisKeys, NamedDims
 using Statistics, OffsetArrays, Tables, UniqueVectors, LazyStack
-using ChainRulesCore: ProjectTo
+using ChainRulesCore: ProjectTo, NoTangent
 using ChainRulesTestUtils: test_rrule
 using FiniteDifferences