Projection for UniformScaling (#533)

* I does his own stunts

* version

Author: mcabbott

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRulesCore"
 uuid = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
-version = "1.11.6"
+version = "1.12"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/projection.jl

@@ -379,6 +379,12 @@ end
 
 using LinearAlgebra: AdjointAbsVec, TransposeAbsVec, AdjOrTransAbsVec
 
+# UniformScaling can represent its own cotangent
+ProjectTo(x::UniformScaling) = ProjectTo{UniformScaling}(; λ=ProjectTo(x.λ))
+ProjectTo(x::UniformScaling{Bool}) = ProjectTo(false)
+(pr::ProjectTo{UniformScaling})(dx::UniformScaling) = UniformScaling(pr.λ(dx.λ))
+(pr::ProjectTo{UniformScaling})(dx::Tangent{<:UniformScaling}) = UniformScaling(pr.λ(dx.λ))
+
 # Row vectors
 ProjectTo(x::AdjointAbsVec) = ProjectTo{Adjoint}(; parent=ProjectTo(parent(x)))
 # Note that while [1 2; 3 4]' isa Adjoint, we use ProjectTo{Adjoint} only to encode AdjointAbsVec.

src/tangent_types/abstract_zero.jl

@@ -37,6 +37,8 @@ Base.sum(z::AbstractZero; dims=:) = z
 Base.reshape(z::AbstractZero, size...) = z
 Base.reverse(z::AbstractZero, args...; kwargs...) = z
 
+(::Type{<:UniformScaling})(z::AbstractZero) = z
+
 """
     ZeroTangent() <: AbstractZero
 

test/projection.jl

@@ -208,6 +208,13 @@ struct NoSuperType end
     ##### `LinearAlgebra`
     #####
 
+    @testset "UniformScaling" begin
+        @test ProjectTo(I)(123) === NoTangent()
+        @test ProjectTo(2 * I)(I * 3im) === 0.0 * I
+        @test ProjectTo((4 + 5im) * I)(Tangent{typeof(im * I)}(; λ = 6)) === (6.0 + 0.0im) * I
+        @test ProjectTo(7 * I)(Tangent{typeof(2I)}()) == ZeroTangent()
+    end
+
     @testset "LinearAlgebra: $adj vectors" for adj in [transpose, adjoint]
         # adjoint vectors
         padj = ProjectTo(adj([1, 2, 3]))