feat: support track numbers via reactant device API (#1533)

Author: avik-pal

lib/MLDataDevices/Project.toml

@@ -1,7 +1,7 @@
 name = "MLDataDevices"
 uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.14.0"
+version = "1.15.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/MLDataDevices/ext/MLDataDevicesReactantExt.jl

@@ -34,9 +34,19 @@ function device_to_kwargs(dev::ReactantDevice, x)
     dev.device === missing || (kwargs = (; kwargs..., device=dev.device))
     if dev.sharding !== missing
         if dev.sharding isa IdDict
-            sharding = (
-                haskey(dev.sharding, x) ? dev.sharding[x] : Reactant.Sharding.NoSharding()
-            )
+            if haskey(dev.sharding, x)
+                sharding = dev.sharding[x]
+            else
+                if all(x -> x isa Reactant.Sharding.NoSharding, values(dev.sharding))
+                    sharding = Reactant.Sharding.NoSharding()
+                else
+                    meshes = unique([
+                        getfield(sharding, :mesh) for sharding in values(dev.sharding)
+                    ])
+                    @assert length(meshes) == 1 "Multiple meshes are not supported."
+                    sharding = Reactant.Sharding.Replicated(only(meshes))
+                end
+            end
             @assert sharding isa Reactant.Sharding.AbstractSharding
             kwargs = (; kwargs..., sharding)
         elseif dev.sharding isa Reactant.Sharding.AbstractSharding
@@ -140,6 +150,12 @@ Profiler.@annotate "Device Transfer (Reactant)" function Adapt.adapt_storage(
 )
     return rng
 end
+Profiler.@annotate "Device Transfer (Reactant)" function Adapt.adapt_storage(
+    dev::ReactantDevice{C,D,S,T,TN}, x::Number
+) where {C,D,S,T,TN}
+    typeof(x) <: TN && return ConcreteRNumber(x; device_to_kwargs(dev, x)...)
+    return x
+end
 
 function Adapt.adapt_storage(
     ::CPUDevice,

lib/MLDataDevices/src/public.jl

@@ -21,20 +21,28 @@ MetalDevice() = MetalDevice{Missing}()
 struct oneAPIDevice{T<:EltypeAdaptorType} <: AbstractGPUDevice end
 oneAPIDevice() = oneAPIDevice{Missing}()
 
-struct ReactantDevice{C,D,S,T<:EltypeAdaptorType} <: AbstractAcceleratorDevice
+struct ReactantDevice{C,D,S,T<:EltypeAdaptorType,TN} <: AbstractAcceleratorDevice
     client::C
     device::D
     sharding::S
 end
 function ReactantDevice()
-    return ReactantDevice{Missing,Missing,Missing,Missing}(missing, missing, missing)
+    return ReactantDevice{Missing,Missing,Missing,Missing,Union{}}(
+        missing, missing, missing
+    )
 end
-function ReactantDevice(client, device, sharding)
-    return ReactantDevice{typeof(client),typeof(device),typeof(sharding),Missing}(
+function ReactantDevice(client, device, sharding, _::Type{TN}=Union{}) where {TN}
+    return ReactantDevice{typeof(client),typeof(device),typeof(sharding),Missing,TN}(
         client, device, sharding
     )
 end
 
+function with_track_numbers(
+    dev::ReactantDevice{C,D,S,T,Union{}}, _::Type{TN}
+) where {C,D,S,T,TN}
+    return ReactantDevice{C,D,S,T,TN}(dev.client, dev.device, dev.sharding)
+end
+
 # Helper functions to get the eltype from device types
 Base.eltype(::CPUDevice{T}) where {T} = T
 Base.eltype(::CUDADevice{D,T}) where {D,T} = T
@@ -70,30 +78,38 @@ function with_eltype(::oneAPIDevice, ::Type{T}) where {T<:AbstractFloat}
     return oneAPIDevice{T}()
 end
 
-function with_eltype(dev::ReactantDevice{C,D,S}, ::Missing) where {C,D,S}
-    return ReactantDevice{C,D,S,Missing}(dev.client, dev.device, dev.sharding)
+function with_eltype(dev::ReactantDevice{C,D,S,<:Any,TN}, ::Missing) where {C,D,S,TN}
+    return ReactantDevice{C,D,S,Missing,TN}(dev.client, dev.device, dev.sharding)
 end
-function with_eltype(dev::ReactantDevice{C,D,S}, ::Nothing) where {C,D,S}
-    return ReactantDevice{C,D,S,Nothing}(dev.client, dev.device, dev.sharding)
+function with_eltype(dev::ReactantDevice{C,D,S,<:Any,TN}, ::Nothing) where {C,D,S,TN}
+    return ReactantDevice{C,D,S,Nothing,TN}(dev.client, dev.device, dev.sharding)
 end
-function with_eltype(dev::ReactantDevice{C,D,S}, ::Type{T}) where {C,D,S,T<:AbstractFloat}
-    return ReactantDevice{C,D,S,T}(dev.client, dev.device, dev.sharding)
+function with_eltype(
+    dev::ReactantDevice{C,D,S,<:Any,TN}, ::Type{T}
+) where {C,D,S,TN,T<:AbstractFloat}
+    return ReactantDevice{C,D,S,T,TN}(dev.client, dev.device, dev.sharding)
 end
 
 function Base.:(==)(
-    x::ReactantDevice{<:Any,<:Any,<:Any,T1}, y::ReactantDevice{<:Any,<:Any,<:Any,T2}
-) where {T1,T2}
+    x::ReactantDevice{<:Any,<:Any,<:Any,T1,TN1}, y::ReactantDevice{<:Any,<:Any,<:Any,T2,TN2}
+) where {T1,T2,TN1,TN2}
     if x.client !== missing && y.client !== missing && x.client.client != y.client.client
         return false
     end
 
-    if x.device !== missing && y.device !== missing && x.device.device != y.device.device
+    if (
+        x.device !== missing &&
+        x.device !== nothing && # can be nothing if objects are sharded
+        y.device !== missing &&
+        y.device !== nothing && # can be nothing if objects are sharded
+        x.device.device != y.device.device
+    )
         return false
     end
 
     T1 === Missing && return T2 === Missing || T2 === Nothing
     T2 === Missing && return T1 === Missing || T1 === Nothing
-    return T1 === T2
+    return T1 === T2 && TN1 === TN2
 end
 
 # XXX: Deprecate in v2
@@ -311,7 +327,8 @@ cpu_device(eltype::T=missing) where {T} = with_eltype(CPUDevice(), eltype)
 
 """
     reactant_device(;
-        force::Bool=false, client=missing, device=missing, sharding=missing, eltype=missing
+        force::Bool=false, client=missing, device=missing, sharding=missing, eltype=missing,
+        track_numbers::Type{TN}=Union{}
     ) -> Union{ReactantDevice, CPUDevice}
 
 Return a `ReactantDevice` object if functional. Otherwise, throw an error if `force` is
@@ -324,19 +341,28 @@ specified, then the default client and index are used.
 `Reactant.Sharding.AbstractSharding` is specified, then we use it to shard all abstract
 arrays. Alternatively, pass in a `IdDict` to specify the sharding for specific leaves.
 
+`track_numbers` can be specified to convert numbers of specified subtypes to be traced.
+
 The `eltype` parameter controls element type conversion:
 
   - `missing/nothing` (default): Preserves the original element type
   - `Type{<:AbstractFloat}`: Converts floating-point arrays to the specified type
 """
 function reactant_device(
-    eltype::T=missing; force::Bool=false, client=missing, device=missing, sharding=missing
-) where {T}
+    eltype::T=missing;
+    force::Bool=false,
+    client=missing,
+    device=missing,
+    sharding=missing,
+    track_numbers::Type{TN}=Union{},
+) where {T,TN}
     msg = "`ReactantDevice` is not loaded or not functional. Load `Reactant.jl` before \
            calling this function. Defaulting to CPU."
     if loaded(ReactantDevice)
         if functional(ReactantDevice)
-            return with_eltype(ReactantDevice(client, device, sharding), eltype)
+            return with_track_numbers(
+                with_eltype(ReactantDevice(client, device, sharding), eltype), track_numbers
+            )
         end
         msg = "`ReactantDevice` is loaded but not functional. Defaulting to CPU."
     end

lib/MLDataDevices/test/reactant_tests.jl

@@ -184,3 +184,13 @@ end
         @test dev(rng) isa Reactant.ReactantRNG
     end
 end
+
+@testset "Track Numbers" begin
+    if MLDataDevices.functional(ReactantDevice)
+        dev = reactant_device(; track_numbers=Float32)
+        x = dev(2.0f0)
+        @test x isa ConcreteRNumber{Float32}
+        x = dev(2.0)
+        @test x isa Float64
+    end
+end