fix ambiguity + oneAPI

LuxDL · Nov 6, 2024 · 9fac54f · 9fac54f
1 parent 71123bd
commit 9fac54f
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 4 deletions.
diff --git a/lib/MLDataDevices/ext/MLDataDevicesCUDAExt.jl b/lib/MLDataDevices/ext/MLDataDevicesCUDAExt.jl
@@ -54,6 +54,7 @@ end
 
 # Device Transfer
 Adapt.adapt_storage(::CUDADevice{Nothing}, x::AbstractArray) = CUDA.cu(x)
+
 function Adapt.adapt_storage(to::CUDADevice, x::AbstractArray)
     old_dev = CUDA.device()  # remember the current device
     dev = MLDataDevices.get_device(x)

diff --git a/lib/MLDataDevices/ext/MLDataDevicesChainRulesExt.jl b/lib/MLDataDevices/ext/MLDataDevicesChainRulesExt.jl
@@ -1,13 +1,19 @@
 module MLDataDevicesChainRulesExt
 
 using Adapt: Adapt
+using MLDataDevices: CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice,OneAPIDevice, ReactantDevice
 using ChainRules: OneElement
-using MLDataDevices: CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, ReactantDevice
 
 Adapt.adapt_storage(::CPUDevice, x::OneElement) = x
-for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, ReactantDevice)
+
+for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, OneAPIDevice, ReactantDevice)
     # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice, ::AbstractArray)
     @eval Adapt.adapt_storage(to::$Dev, x::OneElement) = Adapt.adapt(to, collect(x))
 end
 
+for Dev in (CUDADevice, AMDGPUDevice)
+    # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice{Nothing}, ::AbstractArray)
+    @eval Adapt.adapt_storage(to::$Dev, x::OneElement) = Adapt.adapt(to, collect(x))
+end
+
 end
diff --git a/lib/MLDataDevices/ext/MLDataDevicesZygoteExt.jl b/lib/MLDataDevices/ext/MLDataDevicesZygoteExt.jl
@@ -1,15 +1,20 @@
 module MLDataDevicesZygoteExt
 
 using Adapt: Adapt
-using MLDataDevices: CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, ReactantDevice
+using MLDataDevices: CPUDevice, CUDADevice, AMDGPUDevice, MetalDevice, OneAPIDevice, ReactantDevice
 using Zygote: OneElement
 
 Adapt.adapt_storage(::CPUDevice, x::OneElement) = x
 
-for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, ReactantDevice)
+for Dev in (CUDADevice, AMDGPUDevice, MetalDevice, OneAPIDevice, ReactantDevice)
     # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice, ::AbstractArray)
     @eval Adapt.adapt_storage(to::$Dev, x::OneElement) = Adapt.adapt(to, collect(x))
 end
 
+for Dev in (CUDADevice, AMDGPUDevice)
+    # use `@eval` to avoid ambiguity with adapt_storage(::CUDADevice{Nothing}, ::AbstractArray)
+    @eval Adapt.adapt_storage(to::$Dev{Nothing}, x::OneElement) = Adapt.adapt(to, collect(x))
+end
+
 end