feat: don't collect FillArrays for reactant (#1471)

* feat: don't collect FillArrays for reactant

* Update lib/MLDataDevices/src/internal.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* fix: missing input

* fix: reactant device merging

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: avik-pal

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.11.3"
+version = "1.12.0"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

lib/MLDataDevices/ext/MLDataDevicesFillArraysExt.jl

@@ -1,13 +1,33 @@
 module MLDataDevicesFillArraysExt
 
 using Adapt: Adapt
-using FillArrays: AbstractFill
-using MLDataDevices: CPUDevice, AbstractDevice, Internal
+using FillArrays: AbstractFill, OneElement, Fill, Ones, Zeros
+using MLDataDevices: CPUDevice, ReactantDevice, AbstractDevice, Internal
 
 Adapt.adapt_structure(::CPUDevice, x::AbstractFill) = x
+function Adapt.adapt_structure(dev::CPUDevice, x::Ones{T}) where {T}
+    return Ones{Adapt.adapt(dev, T)}(axes(x))
+end
+function Adapt.adapt_structure(dev::CPUDevice, x::Zeros{T}) where {T}
+    return Zeros{Adapt.adapt(dev, T)}(axes(x))
+end
+Adapt.adapt_structure(dev::CPUDevice, x::Fill) = Fill(Adapt.adapt(dev, x.value), axes(x))
+function Adapt.adapt_structure(dev::CPUDevice, x::OneElement)
+    return OneElement(Adapt.adapt(dev, x.val), x.ind, x.axes)
+end
+
+Adapt.adapt_structure(dev::ReactantDevice, x::AbstractFill) = Internal.to_rarray(dev, x)
+Adapt.adapt_structure(dev::ReactantDevice, x::OneElement) = Internal.to_rarray(dev, x)
+
 Adapt.adapt_structure(to::AbstractDevice, x::AbstractFill) = Adapt.adapt(to, collect(x))
+Adapt.adapt_structure(to::AbstractDevice, x::OneElement) = Adapt.adapt(to, collect(x))
+
+Internal.get_device(::AbstractFill{T}) where {T} = Internal.get_device(T)
+Internal.get_device(f::Fill) = Internal.get_device(f.value)
+Internal.get_device(e::OneElement) = Internal.get_device(e.val)
 
-Internal.get_device(::AbstractFill) = CPUDevice()
-Internal.get_device_type(::AbstractFill) = CPUDevice
+Internal.get_device_type(::AbstractFill{T}) where {T} = Internal.get_device_type(T)
+Internal.get_device_type(f::Fill) = Internal.get_device_type(f.value)
+Internal.get_device_type(e::OneElement) = Internal.get_device_type(e.val)
 
 end

lib/MLDataDevices/ext/MLDataDevicesReactantExt.jl

@@ -27,23 +27,45 @@ function Reactant.make_tracer(
     return prev
 end
 
+# Call into Reactant.to_rarray
+function device_to_kwargs(dev::ReactantDevice, x)
+    kwargs = (;)
+    dev.client === missing || (kwargs = (; kwargs..., client=dev.client))
+    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()
+            )
+            @assert sharding isa Reactant.Sharding.AbstractSharding
+            kwargs = (; kwargs..., sharding)
+        elseif dev.sharding isa Reactant.Sharding.AbstractSharding
+            kwargs = (; kwargs..., dev.sharding)
+        else
+            throw(ArgumentError("`sharding` must be an `IdDict` or a \
+                                 `Reactant.Sharding.AbstractSharding` but got \
+                                 $(typeof(dev.sharding))."))
+        end
+    end
+    return kwargs
+end
+
+function Internal.to_rarray_internal(dev::ReactantDevice, x)
+    return Reactant.to_rarray(x; device_to_kwargs(dev, x)...)
+end
+
 # Default RNG
 MLDataDevices.default_device_rng(::ReactantDevice) = Reactant.TracedRandom.default_rng()
 
 # Query Device from Array
-@static if isdefined(Reactant, :ConcreteIFRTArray)
-    const AllConcreteTypes = Union{
-        Reactant.ConcreteIFRTNumber,
-        Reactant.ConcreteIFRTArray,
-        Reactant.ConcretePJRTNumber,
-        Reactant.ConcretePJRTArray,
-    }
-elseif isdefined(Reactant, :ConcretePJRTArray)
-    const AllConcreteTypes = Union{Reactant.ConcretePJRTNumber,Reactant.ConcretePJRTArray}
-else
-    const AllConcreteTypes = Union{ConcreteRNumber,ConcreteRArray}
-end
+const AllConcreteTypes = Union{
+    <:Reactant.ConcreteIFRTNumber,
+    <:Reactant.ConcreteIFRTArray,
+    <:Reactant.ConcretePJRTNumber,
+    <:Reactant.ConcretePJRTArray,
+}
 
+Internal.get_device(::Type{<:AllConcreteTypes}) = ReactantDevice()
 function Internal.get_device(x::AllConcreteTypes)
     return ReactantDevice(
         Reactant.XLA.client(x),
@@ -53,6 +75,7 @@ function Internal.get_device(x::AllConcreteTypes)
         ),
     )
 end
+Internal.get_device_type(::Type{<:AllConcreteTypes}) = ReactantDevice
 Internal.get_device_type(::AllConcreteTypes) = ReactantDevice
 
 function Internal.get_device(::Union{TracedRArray,TracedRNumber})
@@ -69,26 +92,21 @@ Internal.unsafe_free_internal!(::Type{ReactantDevice}, x::AbstractArray) = nothi
 Profiler.@annotate "Device Transfer (Reactant)" 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))
-    if dev.sharding !== missing
-        if dev.sharding isa IdDict
-            sharding =
-                haskey(dev.sharding, x) ? dev.sharding[x] : Reactant.Sharding.NoSharding()
-            @assert sharding isa Reactant.Sharding.AbstractSharding
-            kwargs = (; kwargs..., sharding)
-        elseif dev.sharding isa Reactant.Sharding.AbstractSharding
-            kwargs = (; kwargs..., dev.sharding)
-        else
-            throw(ArgumentError("`sharding` must be an `IdDict` or a \
-                                 `Reactant.Sharding.AbstractSharding` but got \
-                                 $(typeof(dev.sharding))."))
-        end
-    end
-    return ConcreteRArray(x; kwargs...)
+    return ConcreteRArray(x; device_to_kwargs(dev, x)...)
 end
 
+function Adapt.adapt_storage(
+    ::CPUDevice,
+    T::Type{<:Union{<:Reactant.ConcretePJRTNumber,<:Reactant.ConcreteIFRTNumber}},
+)
+    return Reactant.unwrapped_eltype(T)
+end
+
+function Adapt.adapt_storage(
+    ::CPUDevice, x::Union{<:Reactant.ConcretePJRTNumber,<:Reactant.ConcreteIFRTNumber}
+)
+    return Reactant.unwrapped_eltype(x)(x)
+end
 Adapt.adapt_storage(::CPUDevice, ::Reactant.ReactantRNG) = Random.default_rng()
 
 end

lib/MLDataDevices/src/internal.jl

@@ -149,11 +149,23 @@ combine_devices(::AbstractDevice, dev::ReactantDevice) = dev
 function combine_devices(dev1::ReactantDevice, dev2::ReactantDevice)
     if dev1 == dev2
         # `merge(...)` of `IdDict` constructs a `Dict`
-        sharding = dev1.sharding
-        for (k, v) in dev2.sharding
-            sharding[k] = v
+        if dev1.sharding isa IdDict
+            sharding = dev1.sharding
+            if dev2.sharding isa IdDict
+                for (k, v) in dev2.sharding
+                    sharding[k] = v
+                end
+            end
+        elseif dev2.sharding isa IdDict
+            sharding = dev2.sharding
+        else
+            sharding = missing
         end
-        return ReactantDevice(dev1.client, dev1.device, sharding)
+
+        client = dev1.client === missing ? dev2.client : dev1.client
+        device = dev1.device === missing ? dev2.device : dev1.device
+
+        return ReactantDevice(client, device, sharding)
     end
     throw(ArgumentError("Objects are on different devices: $(dev1) and $(dev2)."))
 end
@@ -211,6 +223,9 @@ for op in (:get_device, :get_device_type)
     end
 end
 
+get_device(::Type{<:Number}) = CPUDevice()
+get_device_type(::Type{<:Number}) = CPUDevice
+
 get_device(_) = UnknownDevice()
 get_device_type(_) = UnknownDevice
 
@@ -263,4 +278,11 @@ end
 
 static_length(t::Tuple) = Val(length(t))
 
+function to_rarray(args...; kwargs...)
+    loaded(ReactantDevice) && return to_rarray_internal(args...; kwargs...)
+    return error("`to_rarray` is only supported with `Reactant` loaded.")
+end
+
+function to_rarray_internal end
+
 end

lib/MLDataDevices/src/public.jl

@@ -16,18 +16,14 @@ struct oneAPIDevice <: AbstractGPUDevice end
 end
 
 function Base.:(==)(x::ReactantDevice, y::ReactantDevice)
-    if x.client !== missing
-        y.client === missing && return false
-        x.client.client != y.client.client && return false
-    else
-        y.client !== missing && return false
+    if x.client !== missing && y.client !== missing && x.client.client != y.client.client
+        return false
     end
-    if x.device !== missing
-        y.device === missing && return false
-        x.device.device != y.device.device && return false
-    else
-        y.device !== missing && return false
+
+    if x.device !== missing && y.device !== missing && x.device.device != y.device.device
+        return false
     end
+
     return true
 end
 

lib/MLDataDevices/test/reactant_tests.jl

@@ -41,12 +41,16 @@ using FillArrays, Zygote  # Extensions
         rng=MersenneTwister(),
         one_elem=Zygote.OneElement(2.0f0, (2, 3), (1:3, 1:4)),
         farray=Fill(1.0f0, (2, 3)),
+        one_elem2=FillArrays.OneElement(2.0f0, (2, 3), (1:3, 1:4)),
+        zeros_fa=Zeros{Float32}((2, 3)),
+        ones_fa=Ones{Float32}((2, 3)),
     )
 
     device = reactant_device()
     aType = MLDataDevices.functional(ReactantDevice) ? Reactant.ConcreteRArray : Array
-    rngType =
+    rngType = (
         MLDataDevices.functional(ReactantDevice) ? Reactant.ReactantRNG : Random.AbstractRNG
+    )
 
     ps_xpu = device(ps)
     @test get_device(ps_xpu) isa ReactantDevice
@@ -70,10 +74,10 @@ using FillArrays, Zygote  # Extensions
 
     if MLDataDevices.functional(ReactantDevice)
         @test ps_xpu.one_elem isa Reactant.RArray
-        @test ps_xpu.farray isa Reactant.RArray
-    else
-        @test ps_xpu.one_elem isa Zygote.OneElement
-        @test ps_xpu.farray isa Fill
+        @test ps_xpu.farray isa Fill{<:Reactant.ConcreteRNumber}
+        @test ps_xpu.one_elem2 isa FillArrays.OneElement{<:Reactant.ConcreteRNumber}
+        @test ps_xpu.zeros_fa isa Zeros{<:Reactant.ConcreteRNumber}
+        @test ps_xpu.ones_fa isa Ones{<:Reactant.ConcreteRNumber}
     end
 
     ps_cpu = cpu_device()(ps_xpu)
@@ -100,10 +104,10 @@ using FillArrays, Zygote  # Extensions
 
     if MLDataDevices.functional(ReactantDevice)
         @test ps_cpu.one_elem isa Array
-        @test ps_cpu.farray isa Array
-    else
-        @test ps_cpu.one_elem isa Zygote.OneElement
-        @test ps_cpu.farray isa Fill
+        @test ps_cpu.farray isa Fill{Float32}
+        @test ps_cpu.one_elem2 isa FillArrays.OneElement{Float32}
+        @test ps_cpu.zeros_fa isa Zeros{Float32}
+        @test ps_cpu.ones_fa isa Ones{Float32}
     end
 
     ps_mixed = (; a=rand(2), b=device(rand(2)))