Repository: invenia/NamedDims.jl (julia)

Issue:
`ProjectTo` is not defined for `NamedDimsArray`
```julia
julia> nda = NamedDimsArray{(:a, :b)}(rand(3, 3))
3×3 NamedDimsArray(::Matrix{Float64}, (:a, :b)):
     → b
↓ a  0.670477  0.9866    0.292266
     0.317949  0.926949  0.76807
     0.99007   0.473338  0.560692

julia> project = ProjectTo(nda)
ProjectTo{AbstractArray}(element = ProjectTo{Float64}(), axes = (Base.OneTo(3), Base.OneTo(3)))

julia> project(rand(3, 3))
3×3 Matrix{Float64}:
 0.0611735  0.684982  0.6573
 0.601941   0.336996  0.782579
 0.477275   0.499212  0.697885
```

whereas it should return a NamedDimsArray.

Pull request:
Closes #200

Patch:
diff --git a/Project.toml b/Project.toml
index 1d4a903..77ab467 100644
--- a/Project.toml
+++ b/Project.toml
@@ -1,7 +1,7 @@
 name = "NamedDims"
 uuid = "356022a1-0364-5f58-8944-0da4b18d706f"
 authors = ["Invenia Technical Computing Corporation"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
diff --git a/src/chainrules.jl b/src/chainrules.jl
index c2012b2..b3d28cc 100644
--- a/src/chainrules.jl
+++ b/src/chainrules.jl
@@ -9,3 +9,13 @@ function ChainRulesCore.rrule(T::Type{<:NamedDimsArray}, values)
     NamedDimsArray_values_pullback(ȳ) = _NamedDimsArray_pullback(ȳ)[1:2]
     return T(values), NamedDimsArray_values_pullback
 end
+
+function ChainRulesCore.ProjectTo(x::NamedDimsArray)
+    return ProjectTo{NamedDimsArray{dimnames(x)}}(; data=ProjectTo(parent(x)))
+end
+
+(project::ProjectTo{NDA})(dx::AbstractZero) where {NDA<:NamedDimsArray} = dx
+function (project::ProjectTo{NDA})(dx) where {NDA<:NamedDimsArray}
+    names = unify_names(dimnames(NDA), dimnames(dx))
+    return NamedDimsArray{names}(project.data(parent(dx)))
+end


Test patch:
diff --git a/test/chainrules.jl b/test/chainrules.jl
index 65edecf..19d1661 100644
--- a/test/chainrules.jl
+++ b/test/chainrules.jl
@@ -1,4 +1,44 @@
 @testset "chainrules.jl" begin
-    test_rrule(NamedDimsArray, rand(2, 3), (:a, :b); check_inferred=VERSION >= v"1.6")
-    test_rrule(NamedDimsArray{(:a, :b)}, (rand(2, 3)); check_inferred=VERSION >= v"1.6")
+    @testset "constructor" begin
+        test_rrule(NamedDimsArray, rand(2, 3), (:a, :b); check_inferred=VERSION >= v"1.6")
+        test_rrule(NamedDimsArray{(:a, :b)}, (rand(2, 3)); check_inferred=VERSION >= v"1.6")
+    end
+
+    @testset "ProjectTo" begin
+        nda = NamedDimsArray{(:a, :b)}(rand(3, 3))
+
+        @testset "NoTangent()" begin
+            @test NoTangent() == ProjectTo(nda)(NoTangent())
+        end
+
+        @testset "(:c, :d) -> (:a, :b) == error" begin
+            ndb = NamedDimsArray{(:c, :d)}(rand(3, 3))
+            @test_throws DimensionMismatch ProjectTo(nda)(ndb)
+        end
+
+        @testset "(:_, :_) -> (:a, :b) == (:a, :b)" begin
+            x = rand(3, 3)
+            projected = @inferred ProjectTo(nda)(x)
+            @test dimnames(projected) == dimnames(nda)
+        end
+
+        @testset "(:a, :_) -> (:a, :b) == (:a, :b)" begin
+            nd1 = NamedDimsArray{(:a, :_)}(rand(3, 3))
+            projected = @inferred ProjectTo(nda)(nd1)
+            @test dimnames(projected) == dimnames(nda)
+        end
+
+        @testset "(:a, :b) -> (:a, :_) == (:a, :b)" begin
+            nd1 = NamedDimsArray{(:a, :_)}(rand(3, 3))
+            projected = @inferred ProjectTo(nd1)(nda) # switched order compared to above
+            @test dimnames(projected) == dimnames(nda)
+        end
+
+        @testset "(:_, :b) -> (:a, :_) == (:a, :b)" begin
+            nd1 = NamedDimsArray{(:a, :_)}(rand(3, 3))
+            nd2 = NamedDimsArray{(:_, :b)}(rand(3, 3))
+            projected = @inferred ProjectTo(nd1)(nd2)
+            @test dimnames(projected) == (:a, :b)
+        end
+    end
 end
diff --git a/test/runtests.jl b/test/runtests.jl
index 187ff99..846b567 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -1,4 +1,5 @@
 using BenchmarkTools
+using ChainRulesCore
 using ChainRulesTestUtils
 using NamedDims
 using SparseArrays