feat: handle RNGs and undef arrays gracefully

Author: avik-pal

.buildkite/pipeline.yml

@@ -1,6 +1,6 @@
 steps:
   - label: "Triggering Pipelines (Pull Request)"
-    if: "build.pull_request.base_branch == 'main'"
+    if: build.branch != "main" && build.tag == null
     agents:
       queue: "juliagpu"
     plugins:

ext/MLDataDevicesAMDGPUExt.jl

@@ -49,8 +49,10 @@ function Internal.get_device(x::AMDGPU.AnyROCArray)
     parent_x === x && return AMDGPUDevice(AMDGPU.device(x))
     return Internal.get_device(parent_x)
 end
+Internal.get_device(::AMDGPU.rocRAND.RNG) = AMDGPUDevice(AMDGPU.device())
 
 Internal.get_device_type(::AMDGPU.AnyROCArray) = AMDGPUDevice
+Internal.get_device_type(::AMDGPU.rocRAND.RNG) = AMDGPUDevice
 
 # Set Device
 function MLDataDevices.set_device!(::Type{AMDGPUDevice}, dev::AMDGPU.HIPDevice)

ext/MLDataDevicesCUDAExt.jl

@@ -29,8 +29,12 @@ function Internal.get_device(x::CUDA.AnyCuArray)
     return MLDataDevices.get_device(parent_x)
 end
 Internal.get_device(x::AbstractCuSparseArray) = CUDADevice(CUDA.device(x.nzVal))
+Internal.get_device(::CUDA.RNG) = CUDADevice(CUDA.device())
+Internal.get_device(::CUDA.CURAND.RNG) = CUDADevice(CUDA.device())
 
 Internal.get_device_type(::Union{<:CUDA.AnyCuArray, <:AbstractCuSparseArray}) = CUDADevice
+Internal.get_device_type(::CUDA.RNG) = CUDADevice
+Internal.get_device_type(::CUDA.CURAND.RNG) = CUDADevice
 
 # Set Device
 MLDataDevices.set_device!(::Type{CUDADevice}, dev::CUDA.CuDevice) = CUDA.device!(dev)

ext/MLDataDevicesGPUArraysExt.jl

@@ -2,9 +2,12 @@ module MLDataDevicesGPUArraysExt
 
 using Adapt: Adapt
 using GPUArrays: GPUArrays
-using MLDataDevices: CPUDevice
+using MLDataDevices: Internal, CPUDevice
 using Random: Random
 
 Adapt.adapt_storage(::CPUDevice, rng::GPUArrays.RNG) = Random.default_rng()
 
+Internal.get_device(rng::GPUArrays.RNG) = Internal.get_device(rng.state)
+Internal.get_device_type(rng::GPUArrays.RNG) = Internal.get_device_type(rng.state)
+
 end

src/internal.jl

@@ -129,15 +129,16 @@ end
 
 for op in (:get_device, :get_device_type)
     cpu_ret_val = op == :get_device ? CPUDevice() : CPUDevice
+    unknown_ret_val = op == :get_device ? UnknownDevice() : UnknownDevice
     not_assigned_msg = "AbstractArray has some undefined references. Giving up, returning \
-                        $(cpu_ret_val)..."
+                        $(unknown_ret_val)..."
 
     @eval begin
         function $(op)(x::AbstractArray{T}) where {T}
             if recursive_array_eltype(T)
                 if any(!isassigned(x, i) for i in eachindex(x))
                     @warn $(not_assigned_msg)
-                    return $(cpu_ret_val)
+                    return $(unknown_ret_val)
                 end
                 return mapreduce(MLDataDevices.$(op), combine_devices, x)
             end

src/public.jl

@@ -232,11 +232,6 @@ const GET_DEVICE_ADMONITIONS = """
 !!! note
 
     Trigger Packages must be loaded for this to return the correct device.
-
-!!! warning
-
-    RNG types currently don't participate in device determination. We will remove this
-    restriction in the future.
 """
 
 # Query Device from Array

test/misc_tests.jl

@@ -154,6 +154,6 @@ end
 @testset "undefined references array" begin
     x = Matrix{Any}(undef, 10, 10)
 
-    @test get_device(x) isa CPUDevice
-    @test get_device_type(x) <: CPUDevice
+    @test get_device(x) isa MLDataDevices.UnknownDevice
+    @test get_device_type(x) <: MLDataDevices.UnknownDevice
 end