fix: adapt to upcoming reactant changes (#1246)

LuxDL · Feb 24, 2025 · 8509073 · 8509073 · avik-pal · Feb 24, 2025
1 parent a979d4f
commit 8509073
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 21 deletions.
diff --git a/lib/MLDataDevices/Project.toml b/lib/MLDataDevices/Project.toml
@@ -1,7 +1,7 @@
 name = "MLDataDevices"
 uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.6.10"
+version = "1.6.11"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/lib/MLDataDevices/ext/MLDataDevicesReactantExt.jl b/lib/MLDataDevices/ext/MLDataDevicesReactantExt.jl
@@ -1,7 +1,7 @@
 module MLDataDevicesReactantExt
 
 using Adapt: Adapt
-using MLDataDevices: MLDataDevices, Internal, ReactantDevice, CPUDevice, get_device_type
+using MLDataDevices: MLDataDevices, Internal, ReactantDevice, CPUDevice
 using Random: Random
 using Reactant: Reactant, XLA, ConcreteRArray, ConcreteRNumber, TracedRArray,
                 TracedRNumber
@@ -32,25 +32,11 @@ end
 Internal.unsafe_free_internal!(::Type{ReactantDevice}, x::AbstractArray) = nothing
 
 # Device Transfer
-function Adapt.adapt_storage(
-        dev::ReactantDevice, x::AbstractArray{<:Reactant.ReactantPrimitive})
-    @warn "ReactantDevice got an array on device: $(get_device_type(x)). We will have to \
-           transfer this via CPU." maxlog=1
-    return Adapt.adapt(dev, Adapt.adapt(CPUDevice(), x))
-end
-
-function Adapt.adapt_storage(dev::ReactantDevice, x::Array{<:Reactant.ReactantPrimitive})
-    client = dev.client === missing ? XLA.default_backend[] : dev.client
-    device = dev.device === missing ? nothing : dev.device
-    return ConcreteRArray(x; client, device)
-end
-
-# XXX: Check for client and device and use faster implementation if possible
-function Adapt.adapt_storage(dev::ReactantDevice, x::ConcreteRArray)
-    dev.client === missing && dev.device === missing && return x
-    @warn "Fetching `client` and `device` from a ConcreteRArray hasn't been implemented \
-           yet. Using slow fallback path." maxlog=1
-    return Adapt.adapt(dev, Adapt.adapt(CPUDevice(), x))
+function Adapt.adapt_storage(dev::ReactantDevice, x::AbstractArray)
+    kwargs = (;)
+    dev.client === missing || (kwargs = (; kwargs..., client=dev.client))
+    dev.device === missing || (kwargs = (; kwargs..., device=dev.device))
+    return ConcreteRArray(x; kwargs...)
 end
 
 Adapt.adapt_storage(::CPUDevice, ::Reactant.ConcreteRNG) = Random.default_rng()

diff --git a/lib/MLDataDevices/src/public.jl b/lib/MLDataDevices/src/public.jl
@@ -9,6 +9,7 @@ end
 struct MetalDevice <: AbstractGPUDevice end
 struct oneAPIDevice <: AbstractGPUDevice end
 
+# TODO: We need to rethink how to store the client and device. what about sharding??
 @kwdef struct ReactantDevice{C, D} <: AbstractAcceleratorDevice
     client::C = missing
     device::D = missing