fix: circular dep due to Adapt

SciML · Nov 15, 2024 · 5f627de · 5f627de
1 parent 2921418
commit 5f627de
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 15 deletions.
diff --git a/Project.toml b/Project.toml
@@ -4,6 +4,7 @@ authors = ["Jonnie Diegelman <[email protected]>"]
 version = "0.15.19"
 
 [deps]
+Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"
@@ -14,7 +15,6 @@ StaticArrayInterface = "0d7ed370-da01-4f52-bd93-41d350b8b718"
 StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c"
 
 [weakdeps]
-Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
@@ -24,7 +24,6 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [extensions]
-ComponentArraysAdaptExt = "Adapt"
 ComponentArraysGPUArraysExt = "GPUArrays"
 ComponentArraysOptimisersExt = "Optimisers"
 ComponentArraysRecursiveArrayToolsExt = "RecursiveArrayTools"

diff --git a/ext/ComponentArraysAdaptExt.jl b/ext/ComponentArraysAdaptExt.jl
diff --git a/src/ComponentArrays.jl b/src/ComponentArrays.jl
@@ -3,6 +3,7 @@ module ComponentArrays
 import ChainRulesCore
 import StaticArrayInterface, ArrayInterface, Functors
 import ConstructionBase
+import Adapt
 
 using LinearAlgebra
 using StaticArraysCore: StaticArray, SArray, SVector, SMatrix

diff --git a/src/componentarray.jl b/src/componentarray.jl
@@ -58,6 +58,14 @@ function ComponentArray(data, ax::AbstractAxis...)
     return LazyArray(ComponentArray(x, axs...) for x in part_data)
 end
 
+function Adapt.adapt_structure(to, x::ComponentArray)
+    data = Adapt.adapt(to, getdata(x))
+    return ComponentArray(data, getaxes(x))
+end
+
+Adapt.adapt_storage(::Type{ComponentArray{T,N,A,Ax}}, xs::AT) where {T,N,A,Ax,AT<:AbstractArray} =
+    Adapt.adapt_storage(A, xs)
+
 # Entry from NamedTuple, Dict, or kwargs
 ComponentArray{T}(nt::NamedTuple) where T = ComponentArray(make_carray_args(T, nt)...)
 ComponentArray{T}(::NamedTuple{(), Tuple{}}) where T = ComponentArray(T[], (FlatAxis(),))