Tracker.data overload (#260)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "ComponentArrays"
 uuid = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 authors = ["Jonnie Diegelman <[email protected]>"]
-version = "0.15.12"
+version = "0.15.13"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

ext/ComponentArraysTrackerExt.jl

@@ -10,6 +10,8 @@ end
 
 Tracker.extract_grad!(ca::ComponentArray) = Tracker.extract_grad!(getdata(ca))
 
+Tracker.data(ca::ComponentArray) = ComponentArray(Tracker.data(getdata(ca)), getaxes(ca))
+
 function Base.materialize(bc::Base.Broadcast.Broadcasted{Tracker.TrackedStyle, Nothing,
     typeof(zero), <:Tuple{<:ComponentVector}})
     ca = first(bc.args)

test/autodiff_tests.jl

@@ -117,3 +117,13 @@ end
 
     @test Δ isa AbstractVector{Float64}
 end
+
+@testset "Tracker untrack" begin
+    ps = Tracker.param(ComponentArray(; a = rand(2)))
+    @test eltype(getdata(ps)) <: Tracker.TrackedReal{Float64}
+
+    ps_data = Tracker.data(ps)
+    @test !(eltype(getdata(ps_data)) <: Tracker.TrackedReal{Float64})
+    @test eltype(getdata(ps_data)) <: Float64
+end
+