Gradient for dot(x,A,y) (#261)

* rrule for 3-arg dot, take 1

* fix complex

* Update test/rulesets/LinearAlgebra/dense.jl

Co-authored-by: willtebbutt <wt0881@my.bristol.ac.uk>

* add method for Diagonal

* restrict to AbstractVector{<:Number} etc

* add PermutedDimsArray tests

* v0.7.30

Co-authored-by: Michael Abbott <me@escbook>
Co-authored-by: willtebbutt <wt0881@my.bristol.ac.uk>

Project.toml

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

src/rulesets/LinearAlgebra/dense.jl

@@ -13,6 +13,35 @@ function rrule(::typeof(dot), x, y)
     return dot(x, y), dot_pullback
 end
 
+function frule((_, Δx, ΔA, Δy), ::typeof(dot), x::AbstractVector{<:Number}, A::AbstractMatrix{<:Number}, y::AbstractVector{<:Number})
+    return dot(x, A, y), dot(Δx, A, y) + dot(x, ΔA, y) + dot(x, A, Δy)
+end
+
+function rrule(::typeof(dot), x::AbstractVector{<:Number}, A::AbstractMatrix{<:Number}, y::AbstractVector{<:Number})
+    Ay = A * y
+    z = adjoint(x) * Ay
+    function dot_pullback(ΔΩ)
+        dx = @thunk conj(ΔΩ) .* Ay
+        dA = @thunk ΔΩ .* x .* adjoint(y)
+        dy = @thunk ΔΩ .* (adjoint(A) * x)
+        return (NO_FIELDS, dx, dA, dy)
+    end
+    dot_pullback(::Zero) = (NO_FIELDS, Zero(), Zero(), Zero())
+    return z, dot_pullback
+end
+
+function rrule(::typeof(dot), x::AbstractVector{<:Number}, A::Diagonal{<:Number}, y::AbstractVector{<:Number})
+    z = dot(x,A,y)
+    function dot_pullback(ΔΩ)
+        dx = @thunk conj(ΔΩ) .* A.diag .* y  # A*y is this broadcast, can be fused
+        dA = @thunk Diagonal(ΔΩ .* x .* conj(y))  # calculate N not N^2 elements
+        dy = @thunk ΔΩ .* conj.(A.diag) .* x
+        return (NO_FIELDS, dx, dA, dy)
+    end
+    dot_pullback(::Zero) = (NO_FIELDS, Zero(), Zero(), Zero())
+    return z, dot_pullback
+end
+
 #####
 ##### `cross`
 #####

test/rulesets/LinearAlgebra/dense.jl

@@ -24,6 +24,23 @@
             frule_test(dot, (x, ẋ), (y, ẏ))
             rrule_test(dot, randn(T), (x, x̄), (y, ȳ))
         end
+        @testset "3-arg dot, Array{$T}" for T in (Float64, ComplexF64)
+            M, N = 3, 4
+            x, A, y = randn(T, M), randn(T, M, N), randn(T, N)
+            ẋ, Adot, ẏ = randn(T, M), randn(T, M, N), randn(T, N)
+            x̄, Abar, ȳ = similar(x), similar(A), similar(y)
+            frule_test(dot, (x, ẋ), (A, Adot), (y, ẏ))
+            rrule_test(dot, randn(T), (x, x̄), (A, Abar), (y, ȳ))
+        end
+        permuteddimsarray(A) = PermutedDimsArray(A, (2,1))
+        @testset "3-arg dot, $F{$T}" for T in (Float32, ComplexF32), F in (adjoint, permuteddimsarray)
+            M, N = 3, 4
+            x, A, y = rand(T, M), F(rand(T, N, M)), rand(T, N)
+            ẋ, Adot, ẏ = rand(T, M), F(rand(T, N, M)), rand(T, N)
+            x̄, Abar, ȳ = similar(x), F(rand(T, N, M)), similar(y)
+            frule_test(dot, (x, ẋ), (A, Adot), (y, ẏ); rtol=1f-3)
+            rrule_test(dot, float(rand(T)), (x, x̄), (A, Abar), (y, ȳ); rtol=1f-3)
+        end
     end
     @testset "cross" begin
         @testset "frule" begin

test/rulesets/LinearAlgebra/structured.jl

@@ -33,7 +33,12 @@
         comp = Composite{typeof(res)}(; diag=10*res.diag)  # this is the structure of Diagonal
         @test pb(comp) == (NO_FIELDS, [10, 40])
     end
-
+    @testset "dot(x, ::Diagonal, y)" begin
+        N = 4
+        x, d, y = randn(ComplexF64, N), randn(ComplexF64, N), randn(ComplexF64, N)
+        D = Diagonal(d)
+        rrule_test(dot, rand(ComplexF64), (x,similar(x)), (D,similar(D)), (y,similar(y)))
+    end
     @testset "::Diagonal * ::AbstractVector" begin
         N = 3
         rrule_test(