Skip to content

Commit

Permalink
fix: restructure of tracked component arrays (#269)
Browse files Browse the repository at this point in the history
  • Loading branch information
@avik-pal
avik-pal authored Sep 10, 2024
1 parent f5574a5 commit fbf6a82
Show file tree
Hide file tree
Showing 4 changed files with 15 additions and 2 deletions.
2 changes: 1 addition & 1 deletion Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "ComponentArrays"
uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
authors = ["Jonnie Diegelman <[email protected]>"]
version = "0.15.16"
version = "0.15.17"

[deps]
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
Expand Down
7 changes: 7 additions & 0 deletions ext/ComponentArraysTrackerExt.jl
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
module ComponentArraysTrackerExt

using ArrayInterface: ArrayInterface
using ComponentArrays, Tracker

function Tracker.param(ca::ComponentArray)
Expand Down Expand Up @@ -34,4 +35,10 @@ end
return ComponentArrays._getindex(Base.getindex, x, v)
end

function ArrayInterface.restructure(x::ComponentVector,
y::ComponentVector{T, <:TrackedArray}) where {T}
getaxes(x) == getaxes(y) || error("Axes must match")
return y
end

end
1 change: 1 addition & 0 deletions test/Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
[deps]
ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
ArrayInterfaceCore = "30b0a656-2188-435a-8636-2ec0e6a096e2"
BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
Expand Down
7 changes: 6 additions & 1 deletion test/autodiff_tests.jl
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,5 @@
import FiniteDiff, ForwardDiff, ReverseDiff, Tracker, Zygote
using Optimisers
using Optimisers, ArrayInterface
using Test

F(a, x) = sum(abs2, a) * x^3
Expand Down Expand Up @@ -127,3 +127,8 @@ end
@test eltype(getdata(ps_data)) <: Float64
end

@testset "ArrayInterface restructure TrackedArray" begin
ps = ComponentArray(; a = rand(2), b = (; c = rand(2)))
ps_tracked = Tracker.param(ps)
@test ArrayInterface.restructure(ps, ps_tracked) isa ComponentVector{<:Any, <:Tracker.TrackedArray}
end

4 comments on commit fbf6a82

@jonniedie
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/114931

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.15.17 -m "<description of version>" fbf6a82130769695eeec50290bd79731d0dfffff
git push origin v0.15.17

@jonniedie
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Error while trying to register: Version 0.15.17 already exists

Please sign in to comment.