add fast paths for arrays of correct numbers, same ndims

Author: mcabbott

src/projection.jl

@@ -40,6 +40,17 @@ Base.propertynames(p::ProjectTo) = propertynames(backing(p))
 backing(project::ProjectTo) = getfield(project, :info)
 
 project_type(p::ProjectTo{T}) where {T} = T
+project_eltype(p::ProjectTo{T}) where {T} = eltype(T)
+
+function project_promote_type(projectors)
+    T = mapreduce(project_type, promote_type, projectors)
+    if T <: Number
+        # The point of this function is to make p.element for arrays. Not in use yet!
+        return ProjectTo(zero(T))
+    else
+        return ProjectTo{Any}()
+    end
+end
 
 function Base.show(io::IO, project::ProjectTo{T}) where {T}
     print(io, "ProjectTo{")
@@ -178,8 +189,16 @@ end
 # no structure worth re-imposing. Then any array is acceptable as a gradient.
 
 # For arrays of numbers, just store one projector:
-function ProjectTo(x::AbstractArray{T}) where {T<:Number}
-    return ProjectTo{AbstractArray}(; element=_eltype_projectto(T), axes=axes(x))
+function ProjectTo(x::AbstractArray{T,N}) where {T<:Number,N}
+    element = _eltype_projectto(T)
+    S = project_type(element)  # new idea -- for any number, S is enough. 
+                               # Store .element for now too, although it's redundant? Reconstruct from eltype?
+    if axes(x) isa NTuple{N,Base.OneTo{Int}}
+        return ProjectTo{AbstractArray{S,N}}(; element=element, axes=axes(x))
+    else
+        # Omitting N prohibits the fast path, and thus won't skip reshape for OffsetArrays, SArrays, etc.
+        return ProjectTo{AbstractArray{S}}(; element=element, axes=axes(x))
+    end
 end
 ProjectTo(x::AbstractArray{Bool}) = ProjectTo{NoTangent}()
 
@@ -197,7 +216,7 @@ function ProjectTo(xs::AbstractArray)
     end
 end
 
-function (project::ProjectTo{AbstractArray})(dx::AbstractArray{S,M}) where {S,M}
+function (project::ProjectTo{<:AbstractArray})(dx::AbstractArray{S,M}) where {S,M}
     # First deal with shape. The rule is that we reshape to add or remove trivial dimensions
     # like dx = ones(4,1), where x = ones(4), but throw an error on dx = ones(1,4) etc.
     dy = if axes(dx) == project.axes
@@ -221,24 +240,33 @@ function (project::ProjectTo{AbstractArray})(dx::AbstractArray{S,M}) where {S,M}
     return dz
 end
 
+# Fast path, for arrays of numbers:
+# (::ProjectTo{AbstractArray{T,N}})(dx::AbstractArray{T,N}) where {T,N} = (@info "fast 1"; dx)
+(::ProjectTo{AbstractArray{T,N}})(dx::AbstractArray{S,N}) where {S<:T} where {T,N} = dx # (@info "fast 2"; dx)
+(project::ProjectTo{AbstractArray{T,N}})(dx::AbstractArray{S,N}) where {S,T,N} = (@info "fast 3"; map(project.element, dx))
+
 # Trivial case, this won't collapse Any[NoTangent(), NoTangent()] but that's OK.
-(project::ProjectTo{AbstractArray})(dx::AbstractArray{<:AbstractZero}) = NoTangent()
+(project::ProjectTo{AbstractArray{T,N}})(dx::AbstractArray{<:AbstractZero}) where {T,N} = NoTangent()
 
 # Row vectors aren't acceptable as gradients for 1-row matrices:
-function (project::ProjectTo{AbstractArray})(dx::LinearAlgebra.AdjOrTransAbsVec)
+# function (project::ProjectTo{<:AbstractArray})(dx::LinearAlgebra.AdjOrTransAbsVec)
+#     return project(reshape(vec(dx), 1, :))
+# end
+function (project::ProjectTo{AbstractArray{T,N}})(dx::LinearAlgebra.AdjOrTransAbsVec) where {T,N}
     return project(reshape(vec(dx), 1, :))
 end
 
 # Zero-dimensional arrays -- these have a habit of going missing,
 # although really Ref() is probably a better structure.
-function (project::ProjectTo{AbstractArray})(dx::Number) # ... so we restore from numbers
-    if !(project.axes isa Tuple{})
-        throw(DimensionMismatch(
-            "array with ndims(x) == $(length(project.axes)) >  0 cannot have dx::Number",
-        ))
-    end
-    return fill(project.element(dx))
-end
+# function (project::ProjectTo{<:AbstractArray})(dx::Number) # ... so we restore from numbers
+#     if !(project.axes isa Tuple{})
+#         throw(DimensionMismatch(
+#             "array with ndims(x) == $(length(project.axes)) >  0 cannot have dx::Number",
+#         ))
+#     end
+#     return fill(project.element(dx))
+# end
+(project::ProjectTo{AbstractArray{<:Number,0}})(dx::Number) = fill(project.element(dx))
 
 # Ref -- works like a zero-array, also allows restoration from a number:
 ProjectTo(x::Ref) = ProjectTo{Ref}(; x=ProjectTo(x[]))

test/projection.jl

@@ -24,7 +24,6 @@ Base.zero(x::Dual) = Dual(zero(x.value), zero(x.partial))
         @test ProjectTo(2.0+3.0im)(1+1im) === 1.0+1.0im
         @test ProjectTo(2.0)(1+1im) === 1.0
 
-
         # storage
         @test ProjectTo(1)(pi) === pi
         @test ProjectTo(1 + im)(pi) === ComplexF64(pi)
@@ -265,6 +264,15 @@ Base.zero(x::Dual) = Dual(zero(x.value), zero(x.partial))
     ##### `OffsetArrays`
     #####
 
+# function ProjectTo(x::OffsetArray{T,N}) where {T<:Number,N}
+#     # As usual:
+#     element = ChainRulesCore._eltype_projectto(T)
+#     S = ChainRulesCore.project_type(element)
+#     # But don't save N? Avoids fast path?
+#     # Or perhaps the default constructor can check whether axes(x) is NTuple{N,OneTo}?
+#     return ProjectTo{AbstractArray{S}}(; element=element, axes=axes(x))
+# end
+
     @testset "OffsetArrays" begin
         # While there is no code for this, the rule that it checks axes(x) == axes(dx) else
         # reshape means that it restores offsets. (It throws an error on nontrivial size mismatch.)