Add ChainRules Overloads to StaticArrays

Author: avik-pal

Project.toml

@@ -1,8 +1,9 @@
 name = "StaticArrays"
 uuid = "90137ffa-7385-5640-81b9-e52037218182"
-version = "1.7.0"
+version = "1.7.1"
 
 [deps]
+ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -17,6 +18,7 @@ StaticArraysStatisticsExt = "Statistics"
 
 [compat]
 Aqua = "0.7"
+ChainRulesCore = "1"
 PrecompileTools = "1"
 StaticArraysCore = "~1.4.0"
 julia = "1.6"

src/StaticArrays.jl

@@ -21,6 +21,10 @@ using LinearAlgebra: checksquare
 
 using PrecompileTools
 
+# ChainRulesCore imports
+import ChainRulesCore: ProjectTo, Tangent, project_type, rrule
+import ChainRulesCore as CRC
+
 # StaticArraysCore imports
 # there is intentionally no "using StaticArraysCore" to not take all symbols exported
 # from StaticArraysCore to make transitioning definitions to StaticArraysCore easier.
@@ -133,6 +137,8 @@ include("flatten.jl")
 include("io.jl")
 include("pinv.jl")
 
+include("chainrules.jl")
+
 @static if !isdefined(Base, :get_extension) # VERSION < v"1.9-"
     include("../ext/StaticArraysStatisticsExt.jl")
 end

src/chainrules.jl

@@ -0,0 +1,23 @@
+# Projecting a tuple to SMatrix leads to ChainRulesCore._projection_mismatch by default, so
+# overloaded here
+function (project::ProjectTo{<:Tangent{<:Tuple}})(dx::StaticArraysCore.SArray)
+    dy = reshape(dx, axes(project.elements))
+    dz = ntuple(i -> project.elements[i](dy[i]), length(project.elements))
+    return project_type(project)(dz...)
+end
+
+# Project SArray to SArray
+function ProjectTo(x::SArray{S, T}) where {S, T}
+    return ProjectTo{SArray}(; element = CRC._eltype_projectto(T), axes = S)
+end
+
+function (project::ProjectTo{SArray})(dx::AbstractArray{S, M}) where {S, M}
+    return SArray{project.axes}(dx)
+end
+
+# Adjoint for SArray constructor
+function rrule(::Type{T}, x::Tuple) where {T <: SArray}
+    project_x = ProjectTo(x)
+    ∇Array(∂y) = (NoTangent(), project_x(∂y))
+    return T(x), ∇Array
+end