diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index b3cb4b73..4fde5205 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -11,27 +11,16 @@ function plus!(out, x, y) return out end -@inline plus!(out::TrackedArray, x::TrackedArray, y::TrackedArray) = record_plus!(out, x, y) - -for A in ARRAY_TYPES - @eval @inline plus!(out::TrackedArray, x::TrackedArray, y::$(A)) = record_plus!(out, x, y) - @eval @inline plus!(out::TrackedArray, x::$(A), y::TrackedArray) = record_plus!(out, x, y) -end - -function record_plus!(out::TrackedArray, x, y) - copyto!(value(out), value(x) + value(y)) - record!(tape(x, y), SpecialInstruction, +, (x, y), out) - return out -end - # Base allocating version @inline Base.:+(x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {X,Y,D} = record_plus(x, y, D) -for A in ARRAY_TYPES - @eval @inline Base.:+(x::TrackedArray{V,D}, y::$(A)) where {V,D} = record_plus(x, y, D) - @eval @inline Base.:+(x::$(A), y::TrackedArray{V,D}) where {V,D} = record_plus(x, y, D) -end +@inline Base.:+(x::TrackedArray{V,D}, y::AbstractArray) where {V,D} = record_plus(x, y, D) +@inline Base.:+(x::AbstractArray, y::TrackedArray{V,D}) where {V,D} = record_plus(x, y, D) + +# avoid ambiguities with StaticArrays +@inline Base.:+(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_plus(x, y, D) +@inline Base.:+(x::StaticArray, y::TrackedArray{V,D}) where {V,D} = record_plus(x, y, D) function record_plus(x, y, ::Type{D}) where D tp = tape(x, y) @@ -78,35 +67,16 @@ function minus!(out, x, y) end end -@inline minus!(out::TrackedArray, x::TrackedArray, y::TrackedArray) = record_minus!(out, x, y) -@inline minus!(out::TrackedArray, x::TrackedArray) = record_minus!(out, x, y) - -for A in ARRAY_TYPES - @eval @inline minus!(out::TrackedArray, x::TrackedArray, y::$(A)) = record_minus!(out, x, y) - @eval @inline minus!(out::TrackedArray, x::$(A), y::TrackedArray) = record_minus!(out, x, y) - @eval @inline minus!(out::TrackedArray, x::$(A)) = record_minus!(out, x) -end - -function record_minus!(out::TrackedArray, x) - copyto!(value(out), -(value(x))) - record!(tape(x), SpecialInstruction, -, x, out) - return out -end - -function record_minus!(out::TrackedArray, x, y) - copyto!(value(out), value(x) - value(y)) - record!(tape(x, y), SpecialInstruction, -, (x, y), out) - return out -end - # Base allocating version -Base.:-(x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {X,Y,D} = record_minus(x, y, D) +@inline Base.:-(x::TrackedArray{X,D}, y::TrackedArray{Y,D}) where {X,Y,D} = record_minus(x, y, D) -for A in ARRAY_TYPES - @eval Base.:-(x::TrackedArray{V,D}, y::$(A)) where {V,D} = record_minus(x, y, D) - @eval Base.:-(x::$(A), y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) -end +@inline Base.:-(x::TrackedArray{V,D}, y::AbstractArray) where {V,D} = record_minus(x, y, D) +@inline Base.:-(x::AbstractArray, y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) + +# avoid ambiguities with StaticArrays +@inline Base.:-(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_minus(x, y, D) +@inline Base.:-(x::StaticArray, y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) function Base.:-(x::TrackedArray{V,D}) where {V,D} tp = tape(x) @@ -273,6 +243,9 @@ else LinearAlgebra.:*(x::Diagonal, y::TrackedVector{V,D}) where {V,D} = record_mul(x, y, D) end +# avoid ambiguities with StaticArrays +LinearAlgebra.:*(x::StaticMatrix, y::TrackedVector{V,D}) where {V,D} = record_mul(x, y, D) + # forward pass # #--------------# diff --git a/src/tracked.jl b/src/tracked.jl index 827e00ea..af140909 100644 --- a/src/tracked.jl +++ b/src/tracked.jl @@ -519,7 +519,16 @@ track(x::AbstractArray, tp::InstructionTape = InstructionTape()) = track(x, elty track(x::Real, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = TrackedReal(x, zero(D), tp) -track(x::AbstractArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} = TrackedArray(x, fill!(similar(x, D), zero(D)), tp) +function track(x::AbstractArray, ::Type{D}, tp::InstructionTape = InstructionTape()) where {D} + v = writable(x) + return TrackedArray(v, fill!(similar(v, D), zero(D)), tp) +end + +# the tape recomputes values in place +writable(x::AbstractArray) = x +writable(x::SArray) = copyto!(similar(x), x) +writable(x::Adjoint) = adjoint(writable(parent(x))) +writable(x::Transpose) = transpose(writable(parent(x))) track!(t::TrackedArray, x::AbstractArray) = (value!(t, x); unseed!(t); t) diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index 415883b8..59bbb833 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/test/derivatives/LinAlgTests.jl @@ -1,6 +1,6 @@ module LinAlgTests -using ReverseDiff, ForwardDiff, Test, LinearAlgebra +using ReverseDiff, ForwardDiff, Test, LinearAlgebra, StaticArrays if get(ENV, "DOWNGRADE_TEST", "false")::String != "true" @test Base.get_extension(ReverseDiff, :StatisticsExt) === nothing @@ -354,4 +354,54 @@ end end end +@testset "`StaticArray`s (#153)" begin + x = rand(2) + X = rand(2, 2) + for (sv, sm) in ( + (SVector(1.0, 2.0), SMatrix{2,2}(1.0, 2.0, 3.0, 4.0)), + (MVector(1.0, 2.0), MMatrix{2,2}(1.0, 2.0, 3.0, 4.0)), + (SizedVector{2}([1.0, 2.0]), SizedMatrix{2,2}([1.0 3.0; 2.0 4.0])), + (SVector{2,BigFloat}(1, 2), SMatrix{2,2,BigFloat}(1, 2, 3, 4)), + ) + for (f, J) in ( + (u -> sv + u, I(2)), + (u -> u + sv, I(2)), + (u -> sv - u, -I(2)), + (u -> u - sv, I(2)), + (u -> sm * u, sm), + ) + test_jacobian(f, x, J) + end + for (f, J) in ( + (U -> sm + U, I(4)), + (U -> U + sm, I(4)), + (U -> sm - U, -I(4)), + (U -> U - sm, I(4)), + (U -> sm * U, kron(I(2), Matrix(sm))), + (U -> U * sm, kron(transpose(Matrix(sm)), I(2))), + (U -> U * sv, kron(transpose(Vector(sv)), I(2))), + ) + test_jacobian(f, X, J) + end + # record at another input so that replaying must recompute the static values + for (f, g) in ( + (u -> sum(abs2, sm * u + sv), u -> 2 * sm' * (sm * u + sv)), + (u -> sum(abs2, u' * sm), u -> 2 * sm * sm' * u), + (u -> sum(abs2, transpose(u) * sm), u -> 2 * sm * transpose(sm) * u), + ) + tp = ReverseDiff.GradientTape(f, rand(2)) + test_approx(ReverseDiff.gradient!(tp, x), g(x)) + end + end + for (f, J) in ( + (u -> SVector(u[1], u[2]) + u, 2I(2)), + (u -> u - SVector(u[1], u[2]), zeros(2, 2)), + (u -> SMatrix{2,2}(u[1], u[2], u[1], u[2]) * u, [2x[1]+x[2] x[1]; x[2] x[1]+2x[2]]), + ) + test_jacobian(f, x, J) + tp = ReverseDiff.JacobianTape(f, rand(2)) + test_approx(ReverseDiff.jacobian!(tp, x), J) + end +end + end # module