fix: preserve object when device is same

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

lib/MLDataDevices/ext/MLDataDevicesAMDGPUExt.jl

@@ -77,7 +77,11 @@ function Internal.unsafe_free_internal!(::Type{AMDGPUDevice}, x::AbstractArray)
 end
 
 # Device Transfer
-Adapt.adapt_storage(::AMDGPUDevice{Nothing}, x::AbstractArray) = AMDGPU.roc(x)
+function Adapt.adapt_storage(::AMDGPUDevice{Nothing}, x::AbstractArray)
+    MLDataDevices.get_device_type(x) <: AMDGPUDevice && return x
+    return AMDGPU.roc(x)
+end
+
 function Adapt.adapt_storage(to::AMDGPUDevice, x::AbstractArray)
     old_dev = AMDGPU.device()  # remember the current device
     dev = MLDataDevices.get_device(x)

lib/MLDataDevices/ext/MLDataDevicesCUDAExt.jl

@@ -57,7 +57,10 @@ function Internal.unsafe_free_internal!(::Type{CUDADevice}, x::AbstractArray)
 end
 
 # Device Transfer
-Adapt.adapt_storage(::CUDADevice{Nothing}, x::AbstractArray) = CUDA.cu(x)
+function Adapt.adapt_storage(::CUDADevice{Nothing}, x::AbstractArray)
+    MLDataDevices.get_device_type(x) <: CUDADevice && return x
+    return CUDA.cu(x)
+end
 
 function Adapt.adapt_storage(to::CUDADevice, x::AbstractArray)
     old_dev = CUDA.device()  # remember the current device

lib/MLDataDevices/ext/MLDataDevicesMetalExt.jl

@@ -29,6 +29,9 @@ function Internal.unsafe_free_internal!(::Type{MetalDevice}, x::AbstractArray)
 end
 
 # Device Transfer
-Adapt.adapt_storage(::MetalDevice, x::AbstractArray) = Metal.mtl(x)
+function Adapt.adapt_storage(::MetalDevice, x::AbstractArray)
+    MLDataDevices.get_device_type(x) <: MetalDevice && return x
+    return Metal.mtl(x)
+end
 
 end

lib/MLDataDevices/ext/MLDataDevicesoneAPIExt.jl

@@ -42,6 +42,7 @@ end
 # Device Transfer
 for (T1, T2) in ((Float64, Float32), (ComplexF64, ComplexF32))
     @eval function Adapt.adapt_storage(::oneAPIDevice, x::AbstractArray{$(T1)})
+        MLDataDevices.get_device_type(x) <: oneAPIDevice && return x
         if !SUPPORTS_FP64[oneAPI.device()]
             @warn LazyString(
                 "Double type is not supported on this device. Using `", $(T2), "` instead.")
@@ -50,6 +51,9 @@ for (T1, T2) in ((Float64, Float32), (ComplexF64, ComplexF32))
         return oneArray(x)
     end
 end
-Adapt.adapt_storage(::oneAPIDevice, x::AbstractArray) = oneArray(x)
+function Adapt.adapt_storage(::oneAPIDevice, x::AbstractArray)
+    MLDataDevices.get_device_type(x) <: oneAPIDevice && return x
+    return oneArray(x)
+end
 
 end

lib/MLDataDevices/src/public.jl

@@ -375,7 +375,10 @@ for op in (:get_device, :get_device_type)
 end
 
 # Adapt Interface
-Adapt.adapt_storage(::CPUDevice, x::AbstractArray) = Array(x)
+function Adapt.adapt_storage(::CPUDevice, x::AbstractArray)
+    get_device_type(x) <: CPUDevice && return x
+    return Array(x)
+end
 Adapt.adapt_storage(to::AbstractDevice, ::Random.TaskLocalRNG) = default_device_rng(to)
 Adapt.adapt_storage(::AbstractDevice, rng::AbstractRNG) = rng