Add rrule for Diagonal * AbstractVector (#108)

* Add rrule for Diagonal * AbstractVector * Simplify rrule definitions for Diagonal * AbstractVector Co-Authored-By: Lyndon White <[email protected]> * Use consistent indentation
JuliaDiff · Sep 21, 2019 · 4afdfdd · 4afdfdd
1 parent 4240832
commit 4afdfdd
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 0 deletions.
diff --git a/src/rulesets/LinearAlgebra/structured.jl b/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(Ȳ)
+        return (NO_FIELDS, @thunk(Diagonal(Ȳ .* V)), @thunk(D * Ȳ))
+    end
+    return D * V, times_pullback
+end
+
 #####
 ##### `Symmetric`
 #####

diff --git a/test/rulesets/LinearAlgebra/structured.jl b/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)))