Skip to content

Support StaticArrays with TrackedArrays (fixes #153) #307

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
devmotion merged 4 commits into from
Sep 29, 2026
Merged

Support StaticArrays with TrackedArrays (fixes #153) #307

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
2ee7e3e
Support `StaticArray`s with `TrackedArray`s (fixes #153)
devmotion Sep 28, 2026
63888c3
Remove unreachable tracked `plus!`/`minus!` methods and extend `Stati…
devmotion Sep 28, 2026
9bbc884
Pin more `StaticArray` combinations and inline `-`
devmotion Sep 29, 2026
4bcaf12
Drop redundant `sm' * u` test and clarify replay comment
devmotion Sep 29, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
59 changes: 16 additions & 43 deletions src/derivatives/linalg/arithmetic.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 #
#--------------#

Expand Down
11 changes: 10 additions & 1 deletion src/tracked.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
52 changes: 51 additions & 1 deletion test/derivatives/LinAlgTests.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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
Expand Down Expand Up @@ -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
Loading