From 2ee7e3eff5c286ef891c00894563abe28220b5a2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 28 Sep 2026 14:54:18 +0200 Subject: [PATCH 1/4] Support `StaticArray`s with `TrackedArray`s (fixes #153) Resolve the `+`/`-`/`*` ambiguities between `StaticArray`s and `TrackedArray`s, and store immutable static values as mutable static arrays so that tapes can recompute them in place. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/derivatives/linalg/arithmetic.jl | 23 ++++++++++++++-------- src/tracked.jl | 11 ++++++++++- test/derivatives/LinAlgTests.jl | 29 +++++++++++++++++++++++++++- 3 files changed, 53 insertions(+), 10 deletions(-) diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index b3cb4b73..3a1443bc 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -28,10 +28,12 @@ end @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) @@ -103,10 +105,12 @@ end 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 +Base.:-(x::TrackedArray{V,D}, y::AbstractArray) where {V,D} = record_minus(x, y, D) +Base.:-(x::AbstractArray, y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) + +# avoid ambiguities with StaticArrays +Base.:-(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_minus(x, y, D) +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 +277,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..a987a56b 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,31 @@ end end end +@testset "`StaticArray`s (#153)" begin + sv = SVector(1.0, 2.0) + sm = SMatrix{2,2}(1.0, 2.0, 3.0, 4.0) + x = rand(2) + 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), + (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]]), + (u -> SVector(u[1], u[2]) .* u[1], [2x[1] 0; x[2] x[1]]), + ) + test_jacobian(f, x, J) + end + # replaying the tape overwrites the static values in place + 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), + ) + tp = ReverseDiff.GradientTape(f, rand(2)) + test_approx(ReverseDiff.gradient!(tp, x), g(x)) + end +end + end # module From 63888c30695b17c3b909f58b22945700dbb475b1 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Mon, 28 Sep 2026 20:42:49 +0200 Subject: [PATCH 2/4] Remove unreachable tracked `plus!`/`minus!` methods and extend `StaticArray` tests The forward passes only call `plus!`/`minus!` on values, which are never `TrackedArray`s, so the methods for tracked outputs could not be reached. One of them also referenced an undefined `y`. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/linalg/arithmetic.jl | 34 ------------------- test/derivatives/LinAlgTests.jl | 50 +++++++++++++++++++--------- 2 files changed, 34 insertions(+), 50 deletions(-) diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index 3a1443bc..90d1a494 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -11,19 +11,6 @@ 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) @@ -80,27 +67,6 @@ 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) diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index a987a56b..4df7741c 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/test/derivatives/LinAlgTests.jl @@ -355,30 +355,48 @@ end end @testset "`StaticArray`s (#153)" begin - sv = SVector(1.0, 2.0) - sm = SMatrix{2,2}(1.0, 2.0, 3.0, 4.0) 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, kron(I(2), Matrix(sm))), + (U -> U * sm, kron(transpose(Matrix(sm)), I(2))), + ) + test_jacobian(f, X, J) + end + # replaying the tape overwrites the static values in place + 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 -> sv + u, I(2)), - (u -> u + sv, I(2)), - (u -> sv - u, -I(2)), - (u -> u - sv, I(2)), - (u -> sm * u, sm), (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]]), - (u -> SVector(u[1], u[2]) .* u[1], [2x[1] 0; x[2] x[1]]), ) test_jacobian(f, x, J) end - # replaying the tape overwrites the static values in place - 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), - ) - tp = ReverseDiff.GradientTape(f, rand(2)) - test_approx(ReverseDiff.gradient!(tp, x), g(x)) - end end end # module From 9bbc884445bc1b2a3745cf57dac48aa2eca40a0c Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 29 Sep 2026 10:35:28 +0200 Subject: [PATCH 3/4] Pin more `StaticArray` combinations and inline `-` Test all orders of `+`/`-` with static matrices, `sm' * u` and `U * sv`, and replay tapes with `TrackedReal`-element static arrays at a new input. Co-Authored-By: Claude Opus 5.5 --- src/derivatives/linalg/arithmetic.jl | 10 +++++----- test/derivatives/LinAlgTests.jl | 6 ++++++ 2 files changed, 11 insertions(+), 5 deletions(-) diff --git a/src/derivatives/linalg/arithmetic.jl b/src/derivatives/linalg/arithmetic.jl index 90d1a494..4fde5205 100644 --- a/src/derivatives/linalg/arithmetic.jl +++ b/src/derivatives/linalg/arithmetic.jl @@ -69,14 +69,14 @@ 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) -Base.:-(x::TrackedArray{V,D}, y::AbstractArray) where {V,D} = record_minus(x, y, D) -Base.:-(x::AbstractArray, y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) +@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 -Base.:-(x::TrackedArray{V,D}, y::StaticArray) where {V,D} = record_minus(x, y, D) -Base.:-(x::StaticArray, y::TrackedArray{V,D}) where {V,D} = record_minus(x, y, D) +@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) diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index 4df7741c..ebd2763e 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/test/derivatives/LinAlgTests.jl @@ -369,14 +369,18 @@ end (u -> sv - u, -I(2)), (u -> u - sv, I(2)), (u -> sm * u, sm), + (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 @@ -396,6 +400,8 @@ end (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 From 4bcaf12a31845a8a9668554d36259f31355826cb Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?David=20M=C3=BCller-Widmann?= Date: Tue, 29 Sep 2026 21:10:12 +0200 Subject: [PATCH 4/4] Drop redundant `sm' * u` test and clarify replay comment StaticArrays' `adjoint` is eager, so `sm' * u` repeated the `sm * u` path. Co-Authored-By: Claude Opus 5.5 --- test/derivatives/LinAlgTests.jl | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/test/derivatives/LinAlgTests.jl b/test/derivatives/LinAlgTests.jl index ebd2763e..59bbb833 100644 --- a/test/derivatives/LinAlgTests.jl +++ b/test/derivatives/LinAlgTests.jl @@ -369,7 +369,6 @@ end (u -> sv - u, -I(2)), (u -> u - sv, I(2)), (u -> sm * u, sm), - (u -> sm' * u, sm'), ) test_jacobian(f, x, J) end @@ -384,7 +383,7 @@ end ) test_jacobian(f, X, J) end - # replaying the tape overwrites the static values in place + # 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),