Add rrule for Diagonal * AbstractVector (#108)

* Add rrule for Diagonal * AbstractVector

* Simplify rrule definitions for Diagonal * AbstractVector

Co-Authored-By: Lyndon White <oxinabox@ucc.asn.au>

* Use consistent indentation

Author: tkf

src/rulesets/LinearAlgebra/structured.jl

@@ -19,6 +19,13 @@ function rrule(::typeof(diag), A::AbstractMatrix)
     return diag(A), diag_pullback
 end
 
+function rrule(::typeof(*), D::Diagonal{<:Real}, V::AbstractVector{<:Real})
+    function times_pullback(Ȳ)
+        return (NO_FIELDS, @thunk(Diagonal(Ȳ .* V)), @thunk(D * Ȳ))
+    end
+    return D * V, times_pullback
+end
+
 #####
 ##### `Symmetric`
 #####

test/rulesets/LinearAlgebra/structured.jl

@@ -7,6 +7,15 @@
         # Concrete type instead of UnionAll
         rrule_test(typeof(D), D, (randn(rng, N), randn(rng, N)))
     end
+    @testset "::Diagonal * ::AbstractVector" begin
+        rng, N = MersenneTwister(123456), 3
+        rrule_test(
+            *,
+            randn(rng, N),
+            (Diagonal(randn(rng, N)), Diagonal(randn(rng, N))),
+            (randn(rng, N), randn(rng, N)),
+        )
+    end
     @testset "diag" begin
         rng, N = MersenneTwister(123456), 7
         rrule_test(diag, randn(rng, N), (randn(rng, N, N), randn(rng, N, N)))