Add rules for sqrt(::Diagonal) (#509)

* WIP add diag sqrt rule

* Adding test with different rand_tangent()

* michaels less allocating suggestion

* spaces and unthunk

* rand the tests

* fix CI

* prevent stochastic failures

Co-authored-by: Miha Zgubic <miha.zgubic@invenialabs.co.uk>

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,19 @@ end
 ##### `Diagonal`
 #####
 
+_diagview(x::Diagonal) = x.diag
+_diagview(x::AbstractMatrix) = view(x, diagind(x))
+_diagview(x::Tangent{<:Diagonal}) = x.diag
+function ChainRulesCore.rrule(::typeof(sqrt), d::Diagonal)
+    y = sqrt(d)
+    @assert y isa Diagonal
+    function sqrt_pullback(Δ)
+        Δ_diag = _diagview(unthunk(Δ))
+        return NoTangent(), Diagonal(Δ_diag ./ (2 .* y.diag))
+    end
+    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(rand(3)))
+        test_rrule(sqrt, Diagonal([1.0, 2]); output_tangent=[1.2 3.4; 1.2 4.3])
+    end
     @testset "$f, $T" for
         f in (Adjoint, adjoint, Transpose, transpose),
         T in (Float64, ComplexF64)