Add rules for sqrt(::Diagonal)

Project.toml

@@ -1,6 +1,6 @@
 name = "ChainRules"
 uuid = "082447d4-558c-5d27-93f4-14fc19e9eca2"
-version = "1.9.0"
+version = "1.10.0"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/rulesets/LinearAlgebra/structured.jl

@@ -35,6 +35,16 @@ end
 ##### `Diagonal`
 #####
 
+_sqrt_pullback(Δ::Tangent, y) = return NoTangent(), Δ / (2y)
+_sqrt_pullback(Δ::Diagonal, y) = return NoTangent(), Δ / (2y)
+_sqrt_pullback(Δ::Matrix, y) = return NoTangent(), Diagonal(Δ) / (2y)
+_sqrt_pullback(Δ::AbstractThunk, y) = return _sqrt_pullback(unthunk(Δ), y)
+function ChainRulesCore.rrule(::typeof(sqrt), d::Diagonal)
+    y = sqrt(d)
+    sqrt_pullback(Δ) = _sqrt_pullback(Δ, y)
+    return y, sqrt_pullback
+end
+
 # these functions are defined outside the rrule because otherwise type inference breaks
 # see https://github.com/JuliaLang/julia/issues/40990
 _Diagonal_pullback(ȳ::AbstractMatrix) = return (NoTangent(), diag(ȳ)) # should we emit a warning here? this shouldn't be called if project works right

test/rulesets/LinearAlgebra/structured.jl

@@ -87,6 +87,10 @@
             end
         end
     end
+    @testset "sqrt(::Diagonal)" begin
+        test_rrule(sqrt, Diagonal([1.0, 2, 3]))
+        test_rrule(sqrt, Diagonal([1.0, 2, 3]); output_tangent=Diagonal(rand(3)) + zeros(3,3))
+    end
     @testset "$f, $T" for
         f in (Adjoint, adjoint, Transpose, transpose),
         T in (Float64, ComplexF64)