use a functor for projection

src/ChainRulesCore.jl

@@ -10,7 +10,7 @@ export RuleConfig, HasReverseMode, NoReverseMode, HasForwardsMode, NoForwardsMod
 export frule_via_ad, rrule_via_ad
 # definition helper macros
 export @non_differentiable, @scalar_rule, @thunk, @not_implemented
-export canonicalize, extern, unthunk  # differential operations
+export ProjectTo, canonicalize, extern, unthunk  # differential operations
 export add!!  # gradient accumulation operations
 # differentials
 export Tangent, NoTangent, InplaceableThunk, Thunk, ZeroTangent, AbstractZero, AbstractThunk
@@ -26,6 +26,7 @@ include("differentials/notimplemented.jl")
 
 include("differential_arithmetic.jl")
 include("accumulation.jl")
+include("projection.jl")
 
 include("config.jl")
 include("rules.jl")

src/differentials/composite.jl

@@ -129,6 +129,7 @@ backing(x::NamedTuple) = x
 backing(x::Dict) = x
 backing(x::Tangent) = getfield(x, :backing)
 
+# For generic structs
 function backing(x::T)::NamedTuple where T
     # note: all computation outside the if @generated happens at runtime.
     # so the first 4 lines of the branchs look the same, but can not be moved out.

src/projection.jl

@@ -0,0 +1,132 @@
+struct ProjectTo{P, D<:NamedTuple}
+    info::D
+end
+ProjectTo{P}(info::D) where {P,D<:NamedTuple} = ProjectTo{P,D}(info)
+ProjectTo{P}(; kwargs...) where {P} = ProjectTo{P}(NamedTuple(kwargs))
+
+backing(project::ProjectTo) = getfield(project, :info)
+Base.getproperty(p::ProjectTo, name::Symbol) = getproperty(backing(p), name)
+Base.propertynames(p::ProjectTo) = propertynames(backing(p))
+
+function Base.show(io::IO, project::ProjectTo{T}) where T
+    print(io, "ProjectTo{")
+    show(io, T)
+    print(io, "}")
+    if isempty(backing(project))
+        print(io, "()")
+    else
+        show(io, backing(project))
+    end
+end
+
+
+"""
+    ProjectTo(x)
+
+Returns a `ProjectTo{P,...}` functor able to project a differential `dx` onto the type `T`
+for a primal `x`.
+This functor encloses over what ever is needed to be able to be able to do that projection.
+For example, when projecting `dx=ZeroTangent()` on an array `P=Array{T, N}`, the size of `x`
+is not available from `P`, so it is stored in the functor.
+"""
+function ProjectTo end
+
+"""
+    (::ProjectTo{T})(dx)
+
+Projects the differential `dx` on the onto type `T`.
+`ProjectTo{T}` is a functor that knows how to perform this projection.
+"""
+function (::ProjectTo) end
+
+# fallback (structs)
+function ProjectTo(x::T) where {T}
+    # Generic fallback for structs, recursively make `ProjectTo`s all their fields
+    #println()
+    #@show x
+    #@show T
+    fields_nt::NamedTuple = backing(x)
+    #@show fields_nt
+    return ProjectTo{T}(map(ProjectTo, fields_nt))
+end
+function (project::ProjectTo{T})(dx::Tangent) where {T}
+    sub_projects = backing(project)
+    #@show sub_projects
+    sub_dxs = backing(canonicalize(dx))
+    #@show sub_dxs
+    _call(f, x) = f(x)
+    return construct(T, map(_call, sub_projects, sub_dxs))
+end
+
+# should not work for Tuples and NamedTuples, as not valid tangent types
+function ProjectTo(x::T) where {T<:Union{<:Tuple, NamedTuple}}
+    throw(ArgumentError("The `x` in `ProjectTo(x)` must be a valid differential, not $x"))
+end
+
+# Generic
+(project::ProjectTo)(dx::AbstractThunk) = project(unthunk(dx))
+(::ProjectTo{T})(dx::T) where {T}  = dx  # not always true, but we can special case for when it isn't
+(::ProjectTo{T})(dx::AbstractZero) where {T} = zero(T)
+
+# Number
+ProjectTo(::T) where {T<:Number} = ProjectTo{T}()
+(::ProjectTo{T})(dx::Number) where {T<:Number} = convert(T, dx)
+(::ProjectTo{T})(dx::Number) where {T<:Real} = convert(T, real(dx))
+
+# Arrays 
+ProjectTo(xs::T) where {T<:Array} = ProjectTo{T}(; elements=map(ProjectTo, xs))
+function (project::ProjectTo{T})(dx::Array) where {T<:Array}
+    _call(f, x) = f(x)
+    return T(map(_call, project.elements, dx))
+end
+function (project::ProjectTo{T})(dx::AbstractZero) where {T<:Array}
+    return T(map(proj->proj(dx), project.elements))
+end
+(project::ProjectTo{<:Array})(dx::AbstractArray) = project(collect(dx))
+
+# Arrays{<:Number}: optimized case so we don't need a projector per element
+ProjectTo(x::T) where {E<:Number, T<:Array{E}} = ProjectTo{T}(; element=ProjectTo(zero(E)), size=size(x))
+(project::ProjectTo{<:Array{T}})(dx::Array) where {T<:Number} = project.element.(dx)
+(project::ProjectTo{<:Array{T}})(dx::AbstractZero) where {T<:Number} = zeros(T, project.size)
+(project::ProjectTo{<:Array{T}})(dx::Tangent{<:SubArray}) where {T<:Number} = project(dx.parent)
+
+# Diagonal
+ProjectTo(x::T) where {T<:Diagonal} = ProjectTo{T}(; diag=ProjectTo(diag(x)))
+(project::ProjectTo{T})(dx::AbstractMatrix) where {T<:Diagonal} = T(project.diag(diag(dx)))
+(project::ProjectTo{T})(dx::AbstractZero) where {T<:Diagonal} = T(project.diag(dx))
+
+# Symmetric and Hermitian
+for SymHerm = (:Symmetric, :Hermitian)
+    @eval begin
+        ProjectTo(x::T) where {T<:$SymHerm} = ProjectTo{T}(; uplo=Symbol(x.uplo), parent=ProjectTo(parent(x)))
+        (project::ProjectTo{<:$SymHerm})(dx::AbstractMatrix) = $SymHerm(project.parent(dx), project.uplo)
+        (project::ProjectTo{<:$SymHerm})(dx::AbstractZero) = $SymHerm(project.parent(dx), project.uplo)
+        (project::ProjectTo{<:$SymHerm})(dx::Tangent) = $SymHerm(project.parent(dx.data), project.uplo)
+    end
+end
+
+# Transpose
+ProjectTo(x::T) where {T<:Transpose} = ProjectTo{T}(; parent=ProjectTo(parent(x)))
+(project::ProjectTo{<:Transpose})(dx::AbstractMatrix) = transpose(project.parent(transpose(dx)))
+(project::ProjectTo{<:Transpose})(dx::Adjoint) = transpose(project.parent(conj(parent(dx))))
+(project::ProjectTo{<:Transpose})(dx::AbstractZero) = transpose(project.parent(dx))
+
+# Adjoint
+ProjectTo(x::T) where {T<:Adjoint} = ProjectTo{T}(; parent=ProjectTo(parent(x)))
+(project::ProjectTo{<:Adjoint})(dx::AbstractMatrix) = adjoint(project.parent(adjoint(dx)))
+(project::ProjectTo{<:Adjoint})(dx::ZeroTangent) = adjoint(project.parent(dx))
+
+# SubArray
+ProjectTo(x::T) where {T<:SubArray} = ProjectTo(copy(x)) # don't project on to a view, but onto matching copy
+
+
+
+
+
+
+
+
+
+
+
+

test/projection.jl

@@ -0,0 +1,215 @@
+struct Fred
+    a::Float64
+end
+Base.zero(::Fred) = Fred(0.0)
+Base.zero(::Type{Fred}) = Fred(0.0)
+
+struct Freddy{T, N}
+    a::Array{T, N}
+end
+Base.:(==)(a::Freddy, b::Freddy) = a.a == b.a
+
+struct Mary
+    a::Fred
+end
+
+struct TwoFields
+    a::Float64
+    b::Float64
+end
+
+@testset "projection" begin
+    @testset "display" begin
+        @test startswith(repr(ProjectTo(Fred(1.1))), "ProjectTo{Fred}(")
+        @test repr(ProjectTo(1.1)) == "ProjectTo{Float64}()"
+    end
+
+    @testset "fallback" begin
+        @test Fred(1.2) == ProjectTo(Fred(1.1))(Fred(1.2))
+        @test Fred(0.0) == ProjectTo(Fred(1.1))(ZeroTangent())
+        @test Fred(3.2) == ProjectTo(Fred(1.1))(@thunk(Fred(3.2)))
+        @test Fred(1.2) == ProjectTo(Fred(1.1))(Tangent{Fred}(;a=1.2))
+
+        # struct with complicated field
+        x = Freddy(zeros(2,2))
+        dx = Tangent{Freddy}(; a=ZeroTangent())
+        @test x == ProjectTo(x)(dx)
+
+        # nested structs
+        f = Fred(0.0)
+        tf = Tangent{Fred}(;a=ZeroTangent())
+        m = Mary(f)
+        dm = Tangent{Mary}(;a=tf)
+        @test m == ProjectTo(m)(dm)
+
+        # two fields
+        tfa = TwoFields(3.0, 0.0)
+        tfb = TwoFields(0.0, 3.0)
+        @test tfa == ProjectTo(tfa)(Tangent{TwoFields}(; a=3.0))
+        @test tfb == ProjectTo(tfb)(Tangent{TwoFields}(; b=3.0))
+    end
+
+    @testset "to Real" begin
+        # Float64
+        @test 3.2 == ProjectTo(1.0)(3.2)
+        @test 0.0 == ProjectTo(1.0)(ZeroTangent())
+        @test 3.2 == ProjectTo(1.0)(@thunk(3.2))
+
+        # down
+        @test 3.2 == ProjectTo(1.0)(3.2 + 3im)
+        @test 3.2f0 == ProjectTo(1.0f0)(3.2)
+        @test 3.2f0 == ProjectTo(1.0f0)(3.2 - 3im)
+
+        # up
+        @test 2.0 == ProjectTo(1.0)(2.0f0)
+    end
+
+    @testset "to Number" begin
+        # Complex
+        @test 2.0 + 4.0im == ProjectTo(1.0im)(2.0 + 4.0im)
+
+        # down
+        @test 2.0 + 0.0im == ProjectTo(1.0im)(2.0)
+        @test 0.0 + 0.0im == ProjectTo(1.0im)(ZeroTangent())
+        @test 0.0 + 0.0im == ProjectTo(1.0im)(@thunk(ZeroTangent()))
+
+        # up
+        @test 2.0 + 0.0im == ProjectTo(1.0im)(2.0)
+    end
+
+    @testset "to Array" begin
+        # to an array of numbers
+        x = zeros(2, 2)
+        @test [1.0 2.0; 3.0 4.0] == ProjectTo(x)([1.0 2.0; 3.0 4.0])
+        @test x == ProjectTo(x)(ZeroTangent())
+
+        x = zeros(2)
+        @test x == ProjectTo(x)(@thunk(ZeroTangent()))
+
+        x = zeros(Float32, 2, 2)
+        @test x == ProjectTo(x)([0.0 0; 0 0])
+
+        x = [1.0 0; 0 4]
+        @test x == ProjectTo(x)(Diagonal([1.0, 4]))
+
+        # to a array of structs
+        x = [Fred(0.0), Fred(0.0)]
+        @test x == ProjectTo(x)([Fred(0.0), Fred(0.0)])
+        @test x == ProjectTo(x)([ZeroTangent(), ZeroTangent()])
+        @test x == ProjectTo(x)([ZeroTangent(), @thunk(Fred(0.0))])
+        @test x == ProjectTo(x)(ZeroTangent())
+        @test x == ProjectTo(x)(@thunk(ZeroTangent()))
+
+        x = [Fred(1.0) Fred(0.0); Fred(0.0) Fred(4.0)]
+        @test x == ProjectTo(x)(Diagonal([Fred(1.0), Fred(4.0)]))
+    end
+
+    @testset "To Array of Arrays" begin
+        # inner arrays have same type but different sizes
+        x = [[1.0, 2.0, 3.0], [4.0, 5.0]]
+        @test x == ProjectTo(x)(x)
+        @test x == ProjectTo(x)([[1.0 + 2im, 2.0, 3.0], [4.0 + 2im, 5.0]])
+
+        # This makes sure we don't fall for https://github.com/JuliaLang/julia/issues/38064
+        @test [[0.0, 0.0, 0.0], [0.0, 0.0]] == ProjectTo(x)(ZeroTangent())
+    end
+
+    @testset "Array{Any} with really messy contents" begin
+        # inner arrays have same type but different sizes
+        x = [[1.0, 2.0, 3.0], [4.0+im 5.0], [[[Fred(1)]]]]
+        @test x == ProjectTo(x)(x)
+        @test x == ProjectTo(x)([[1.0+im, 2.0, 3.0], [4.0+im 5.0], [[[Fred(1)]]]])
+        # using a different type for the 2nd element (Adjoint)
+        @test x == ProjectTo(x)([[1.0+im, 2.0, 3.0], [4.0-im, 5.0]', [[[Fred(1)]]]])
+
+        @test [[0.0, 0.0, 0.0], [0.0im 0.0], [[[Fred(0)]]]] == ProjectTo(x)(ZeroTangent())
+    end
+
+    @testset "to Diagonal" begin
+        d_F64 = Diagonal([0.0, 0.0])
+        d_F32 = Diagonal([0.0f0, 0.0f0])
+        d_C64 = Diagonal([0.0 + 0im, 0.0])
+        d_Fred = Diagonal([Fred(0.0), Fred(0.0)])
+
+        # from Matrix
+        @test d_F64 == ProjectTo(d_F64)(zeros(2, 2))
+        @test d_F64 == ProjectTo(d_F64)(zeros(Float32, 2, 2))
+        @test d_F64 == ProjectTo(d_F64)(zeros(ComplexF64, 2, 2))
+
+        # from Diagonal of Numbers
+        @test d_F64 == ProjectTo(d_F64)(d_F64)
+        @test d_F64 == ProjectTo(d_F64)(d_F32)
+        @test d_F64 == ProjectTo(d_F64)(d_C64)
+
+        # from Diagonal of AbstractTangent
+        @test d_F64 == ProjectTo(d_F64)(ZeroTangent())
+        @test d_C64 == ProjectTo(d_C64)(ZeroTangent())
+        @test d_F64 == ProjectTo(d_F64)(@thunk(ZeroTangent()))
+        @test d_F64 == ProjectTo(d_F64)(Diagonal([ZeroTangent(), ZeroTangent()]))
+        @test d_F64 == ProjectTo(d_F64)(Diagonal([ZeroTangent(), @thunk(ZeroTangent())]))
+
+        # from Diagonal of structs
+        @test d_Fred == ProjectTo(d_Fred)(ZeroTangent())
+        @test d_Fred == ProjectTo(d_Fred)(@thunk(ZeroTangent()))
+        @test d_Fred == ProjectTo(d_Fred)(Diagonal([ZeroTangent(), ZeroTangent()]))
+
+        # from Tangent
+        @test d_F64 == ProjectTo(d_F64)(Tangent{Diagonal}(;diag=[0.0, 0.0]))
+        @test d_F64 == ProjectTo(d_F64)(Tangent{Diagonal}(;diag=[0.0f0, 0.0f0]))
+        @test d_F64 == ProjectTo(d_F64)(Tangent{Diagonal}(;diag=[ZeroTangent(), @thunk(ZeroTangent())]))
+    end
+
+    @testset "to $SymHerm" for SymHerm in (Symmetric, Hermitian)
+        data = [1.0+1im 2-2im; 3 4]
+
+        x = SymHerm(data)
+        @test x == ProjectTo(x)(data)
+        @test x == ProjectTo(x)(Tangent{typeof(x)}(; data=data, uplo=NoTangent()))
+
+        x = SymHerm(data, :L)
+        @test x == ProjectTo(x)(data)
+
+        data = [1.0-2im 0; 0 4]
+        x = SymHerm(data)
+        @test x == ProjectTo(x)(Diagonal([1.0-2im, 4.0]))
+
+        data = [0.0+0im 0; 0 0]
+        x = SymHerm(data)
+        @test x == ProjectTo(x)(ZeroTangent())
+        @test x == ProjectTo(x)(@thunk(ZeroTangent()))
+    end
+
+    @testset "to Transpose" begin
+        x = rand(ComplexF64, 3, 4)
+        t = transpose(x)
+        mt = collect(t)
+        a = adjoint(x)
+        ma = collect(a)
+
+        @test t == ProjectTo(t)(mt)
+        @test conj(t) == ProjectTo(t)(ma)
+        @test zeros(4, 3) == ProjectTo(t)(ZeroTangent())
+        @test zeros(4, 3) == ProjectTo(t)(Tangent{Transpose}(; parent=ZeroTangent()))
+    end
+
+    @testset "to Adjoint" begin
+        x = rand(ComplexF64, 3, 4)
+        a = adjoint(x)
+        ma = collect(a)
+
+        @test a == ProjectTo(a)(ma)
+        @test zeros(4, 3) == ProjectTo(a)(ZeroTangent())
+        @test zeros(4, 3) == ProjectTo(a)(Tangent{Adjoint}(; parent=ZeroTangent()))
+    end
+
+    @testset "to SubArray" begin
+        x = rand(3, 4)
+        sa = view(x, :, 1:2)
+        m = collect(sa)
+
+        # make sure it converts the view to the parent type
+        @test ProjectTo(sa)(m) isa Matrix
+        @test ProjectTo(sa)(ZeroTangent())
+        @test ProjectTo(sa)(Tangent{SubArray}(; parent=ZeroTangent())) isa Matrix
+    end
+end

test/runtests.jl

@@ -16,6 +16,7 @@ using Test
     end
 
     include("accumulation.jl")
+    include("projection.jl")
 
     include("rules.jl")
     include("rule_definition_tools.jl")