diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index 56e78c3f..b3cb4b73 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/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 diff --git a/test/Project.toml b/test/Project.toml index 33298a9e..80006f85 100644 --- a/test/Project.toml +++ b/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" diff --git a/test/api/JacobianTests.jl b/test/api/JacobianTests.jl index 0f26db34..ed3e68c2 100644 --- a/test/api/JacobianTests.jl +++ b/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)) diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index dd53bdfa..014d8c77 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/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