Fix derivatives w.r.t. structured matrices and require DiffTests 0.1.3

src/derivatives/linalg/arithmetic.jl

@@ -165,6 +165,27 @@ mulargpullvalue!(x) = pull_value!(x)
 mulargpullvalue!(x::Adjoint) = pull_value!(adjoint(x))
 mulargpullvalue!(x::Transpose) = pull_value!(transpose(x))
 
+# derivatives w.r.t. structured matrices retain their structure
+function structured_mul!(
+    C::AbstractVecOrMat,
+    A::Union{Number,AbstractVecOrMat},
+    B::Union{Number,AbstractVecOrMat},
+)
+    return mul!(C, A, B)
+end
+function structured_mul!(C::Diagonal, A::AbstractVecOrMat, B::AbstractVecOrMat)
+    map!(dot, C.diag, eachcol(adjoint(A)), eachcol(B))
+    return C
+end
+function structured_mul!(
+    C::Union{UpperTriangular,LowerTriangular,UnitUpperTriangular,UnitLowerTriangular},
+    A::AbstractVecOrMat,
+    B::AbstractVecOrMat,
+)
+    mul!(parent(C), A, B)
+    return C
+end
+
 # recording pass #
 #----------------#
 
@@ -239,6 +260,19 @@ for S1 in (:TrackedArray, :TrackedVector, :TrackedMatrix)
     end
 end
 
+# avoid ambiguities with LinearAlgebra's methods for structured matrices
+@static if VERSION < v"1.11"
+    for W in (:Diagonal, :(LinearAlgebra.AbstractTriangular))
+        @eval begin
+            LinearAlgebra.:*(x::$W, y::TrackedVector{V,D}) where {V,D} = record_mul(x, y, D)
+            LinearAlgebra.:*(x::$W, y::TrackedMatrix{V,D}) where {V,D} = record_mul(x, y, D)
+            LinearAlgebra.:*(x::TrackedMatrix{V,D}, y::$W) where {V,D} = record_mul(x, y, D)
+        end
+    end
+else
+    LinearAlgebra.:*(x::Diagonal, y::TrackedVector{V,D}) where {V,D} = record_mul(x, y, D)
+end
+
 # forward pass #
 #--------------#
 
@@ -276,66 +310,12 @@ function reverse_mul!(output, output_deriv, a, b, a_tmp, b_tmp)
             # involve outer-products of vectors, for such functions, the target
             # a_temp is a vector, but when b is a matrix, we cannot multiply into a vector,
             # so need to reshape memory to look like matrix (see PositiveFactorizations.jl)
-            increment_deriv!(a, mul!(reshape(a_tmp, :, 1), output_deriv, transpose(value(b))))
+            increment_deriv!(a, structured_mul!(reshape(a_tmp, :, 1), output_deriv, transpose(value(b))))
         else
-            increment_deriv!(a, mul!(a_tmp, output_deriv, transpose(value(b))))
+            increment_deriv!(a, structured_mul!(a_tmp, output_deriv, transpose(value(b))))
         end
     end
-    istracked(b) && increment_deriv!(b, mul!(b_tmp, transpose(value(a)), output_deriv))
-end
-
-for (f, F) in ((:transpose, :Transpose), (:adjoint, :Adjoint))
-    @eval begin
-        # a * f(b)
-        function reverse_mul!(output, output_deriv, a, b::$F, a_tmp, b_tmp)
-            _b = ($f)(b)
-            if istracked(a)
-                if a_tmp isa AbstractVector
-                    increment_deriv!(a, mul!(reshape(a_tmp, :, 1), output_deriv, mulargvalue(_b)))
-                else
-                    increment_deriv!(a, mul!(a_tmp, output_deriv, mulargvalue(b)))
-                end
-            end
-            istracked(_b) && increment_deriv!(_b, ($f)(mul!(($f)(b_tmp), ($f)(output_deriv), value(a))))
-        end
-           # f(a) * b
-        function reverse_mul!(output, output_deriv, a::$F, b, a_tmp, b_tmp)
-            _a = ($f)(a)
-            istracked(_a) && increment_deriv!(_a, ($f)(mul!(a_tmp, output_deriv, ($f)(value(b)))))
-            istracked(b) && increment_deriv!(b, mul!(b_tmp, ($f)(mulargvalue(a)), ($f)(output_deriv)))
-        end
-        # f(a) * f(b)
-        function reverse_mul!(output, output_deriv, a::$F, b::$F, a_tmp, b_tmp)
-            _a = ($f)(a)
-            _b = ($f)(b)
-            istracked(_a) && increment_deriv!(_a, ($f)(mul!(a_tmp, ($f)(mulargvalue(b)), ($f)(output_deriv))))
-            istracked(_b) && increment_deriv!(_b, ($f)(mul!(b_tmp, ($f)(output_deriv), ($f)(mulargvalue(a)))))
-        end
-    end
-end
-
-# adjoint(a) * transpose(b)
-
-function reverse_mul!(output, output_deriv, a::Adjoint, b::Transpose, a_tmp, b_tmp)
-    _a = adjoint(a)
-    _b = transpose(b)
-    if istracked(_a)
-        reverse_mul!(output, output_deriv, transpose(_a), b, a_tmp, b_tmp)
-    elseif istracked(_b)
-        increment_deriv!(_b, transpose(mul!(b_tmp, adjoint(output_deriv), adjoint(mulargvalue(a)))))
-    end
-end
-
-# transpose(a) * adjoint(b)
-
-function reverse_mul!(output, output_deriv, a::Transpose, b::Adjoint, a_tmp, b_tmp)
-    _a = transpose(a)
-    _b = adjoint(b)
-    if istracked(_b)
-        reverse_mul!(output, output_deriv, a, transpose(_b), a_tmp, b_tmp)
-    elseif istracked(_a)
-        increment_deriv!(_a, transpose(mul!(a_tmp, adjoint(mulargvalue(b)), adjoint(output_deriv))))
-    end
+    istracked(b) && increment_deriv!(b, structured_mul!(b_tmp, transpose(value(a)), output_deriv))
 end
 
 ## zero

test/Project.toml

@@ -20,7 +20,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 ChainRulesCore = "1"
 DiffResults = "1"
 DiffRules = "1.4"
-DiffTests = "< 0.1.2"
+DiffTests = "0.1.3"
 FillArrays = "1"
 ForwardDiff = "1"
 IrrationalConstants = "0.2"

test/api/JacobianTests.jl

@@ -251,6 +251,11 @@ for f in (DiffTests.ARRAY_TO_ARRAY_FUNCS..., DiffTests.MATRIX_TO_MATRIX_FUNCS...
     test_unary_jacobian(f, rand(5, 5))
 end
 
+for f in DiffTests.VECTOR_TO_VECTOR_FUNCS
+    test_println("VECTOR_TO_VECTOR_FUNCS", f)
+    test_unary_jacobian(f, rand(5))
+end
+
 for f! in DiffTests.INPLACE_ARRAY_TO_ARRAY_FUNCS
     test_println("INPLACE_ARRAY_TO_ARRAY_FUNCS", f!)
     test_unary_jacobian(f!, rand(25), rand(25))

test/derivatives/LinAlgTests.jl

@@ -289,4 +289,54 @@ test_arr2arr_inplace(mul!, *, x, transpose(a), adjoint(b), tp)
 A = [1 2; 3 4]; x = [5, 6];
 @test ReverseDiff.gradient(y -> sum(y'*A), x) == [3, 7]
 
+# Derivatives w.r.t. structured factors (https://github.com/JuliaDiff/DiffTests.jl/issues/13)
+x = rand(3, 3)
+# For a structural projection `S`, d(S(Y) * Y) = S(dY) * Y + S(Y) * dY and
+# d(Y * S(Y)) = dY * S(Y) + Y * S(dY)
+mask(S) = Diagonal(vec(Matrix(S(ones(3, 3)))))
+jac_lmul(S) = kron(transpose(x), I(3)) * mask(S) + kron(I(3), Matrix(S(x)))
+jac_rmul(S) = kron(transpose(Matrix(S(x))), I(3)) + kron(I(3), x) * mask(S)
+for (f, J) in (
+    (y -> Diagonal(y) * y, jac_lmul(Diagonal)),
+    (y -> UpperTriangular(y) * y, jac_lmul(UpperTriangular)),
+    (y -> y * LowerTriangular(y), jac_rmul(LowerTriangular)),
+    (y -> transpose(UpperTriangular(transpose(y))) * y, jac_lmul(LowerTriangular)),
+    (y -> y * adjoint(LowerTriangular(adjoint(y))), jac_rmul(UpperTriangular)),
+)
+    test_println("*(A, B) functions with structured matrices", f)
+    test_approx(ReverseDiff.jacobian(f, x), J)
+    ctp = ReverseDiff.compile(ReverseDiff.JacobianTape(f, x))
+    test_approx(ReverseDiff.jacobian!(ctp, x), J)
+end
+
+function test_jacobian(f, x, J)
+    test_approx(ReverseDiff.jacobian(f, x), J)
+    ctp = ReverseDiff.compile(ReverseDiff.JacobianTape(f, x))
+    test_approx(ReverseDiff.jacobian!(ctp, x), J)
+end
+
+# Constant structured factors
+c = rand(3, 3)
+for W in (Diagonal(diag(c)), UpperTriangular(c), LowerTriangular(c), UnitUpperTriangular(c), UnitLowerTriangular(c))
+    test_println("*(A, B) functions with structured matrices", typeof(W))
+    test_jacobian(y -> W * y, rand(3), Matrix(W))
+    test_jacobian(y -> W * y, rand(3, 2), kron(I(2), Matrix(W)))
+    test_jacobian(y -> y * W, rand(2, 3), kron(transpose(Matrix(W)), I(2)))
+end
+
+# Row vectors
+v, V, M = rand(3), rand(2, 1), rand(3, 2)
+for f in (transpose, adjoint)
+    test_println("*(A, B) functions with row vectors", f)
+    for (h, x) in (
+        ((f, u) -> u * f(u), rand(3)),
+        ((f, u) -> u * f(v), rand(2, 1)),
+        ((f, u) -> V * f(u), rand(3)),
+        ((f, u) -> f(u) * M, rand(3)),
+        ((f, u) -> f(v) * u, rand(3, 2)),
+    )
+        test_jacobian(u -> h(f, u), x, ReverseDiff.jacobian(u -> h(x -> reshape(x, 1, :), u), x))
+    end
+end
+
 end # module