Gradient for dot(x,A,y)

src/rulesets/LinearAlgebra/dense.jl

@@ -19,6 +19,23 @@ function rrule(::typeof(dot), x, y)
     return dot(x, y), dot_pullback
 end
 
+function frule((_, Δx, ΔA, Δy), ::typeof(dot), x::AbstractVector, A::AbstractMatrix, y::AbstractVector)
+    return dot(x, A, y), dot(Δx, A, y) + dot(x, ΔA, y) + dot(x, A, Δy)
+end
+
+function rrule(::typeof(dot), x::AbstractVector, A::AbstractMatrix, y::AbstractVector)
+    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
+
 #####
 ##### `cross`
 #####

test/rulesets/LinearAlgebra/dense.jl

@@ -24,6 +24,14 @@
             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, ȳ = randn(T, M), randn(T, M,N), randn(T, N)
+            frule_test(dot, (x, ẋ), (A, Adot), (y, ẏ))
+            rrule_test(dot, randn(T), (x, x̄), (A, Abar), (y, ȳ))
+        end
     end
     @testset "cross" begin
         @testset "frule" begin