fix: unsafe free for OneHotArrays (#1067)

* fix: unsafe free for OneHotArrays * refactor: warn on un-implemented operations
LuxDL · Nov 11, 2024 · 44750b6 · 44750b6 · avik-pal · Nov 11, 2024
1 parent 62c930b
commit 44750b6
Show file tree

Hide file tree

Showing 7 changed files with 36 additions and 9 deletions.
diff --git a/examples/HyperNet/main.jl b/examples/HyperNet/main.jl
@@ -17,7 +17,8 @@ function load_dataset(::Type{dset}, n_train::Int, n_eval::Int, batchsize::Int) w
 
     return (
         DataLoader((x_train, y_train); batchsize=min(batchsize, n_train), shuffle=true),
-        DataLoader((x_test, y_test); batchsize=min(batchsize, n_eval), shuffle=false))
+        DataLoader((x_test, y_test); batchsize=min(batchsize, n_eval), shuffle=false)
+    )
 end
 
 function load_datasets(n_train=1024, n_eval=32, batchsize=256)
@@ -50,8 +51,11 @@ end
 function create_model()
     ## Doesn't need to be a MLP can have any Lux Layer
     core_network = Chain(FlattenLayer(), Dense(784, 256, relu), Dense(256, 10))
-    weight_generator = Chain(Embedding(2 => 32), Dense(32, 64, relu),
-        Dense(64, Lux.parameterlength(core_network)))
+    weight_generator = Chain(
+        Embedding(2 => 32),
+        Dense(32, 64, relu),
+        Dense(64, Lux.parameterlength(core_network))
+    )
 
     model = HyperNet(weight_generator, core_network)
     return model

diff --git a/lib/MLDataDevices/Project.toml b/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.5.2"
+version = "1.5.3"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/lib/MLDataDevices/ext/MLDataDevicesAMDGPUExt.jl b/lib/MLDataDevices/ext/MLDataDevicesAMDGPUExt.jl
@@ -68,7 +68,11 @@ end
 
 # unsafe_free!
 function Internal.unsafe_free_internal!(::Type{AMDGPUDevice}, x::AbstractArray)
-    AMDGPU.unsafe_free!(x)
+    if applicable(AMDGPU.unsafe_free!, x)
+        AMDGPU.unsafe_free!(x)
+    else
+        @warn "AMDGPU.unsafe_free! is not defined for $(typeof(x))." maxlog=1
+    end
     return
 end
 

diff --git a/lib/MLDataDevices/ext/MLDataDevicesCUDAExt.jl b/lib/MLDataDevices/ext/MLDataDevicesCUDAExt.jl
@@ -48,7 +48,11 @@ end
 
 # unsafe_free!
 function Internal.unsafe_free_internal!(::Type{CUDADevice}, x::AbstractArray)
-    CUDA.unsafe_free!(x)
+    if applicable(CUDA.unsafe_free!, x)
+        CUDA.unsafe_free!(x)
+    else
+        @warn "CUDA.unsafe_free! is not defined for $(typeof(x))." maxlog=1
+    end
     return
 end
 

diff --git a/lib/MLDataDevices/ext/MLDataDevicesMetalExt.jl b/lib/MLDataDevices/ext/MLDataDevicesMetalExt.jl
@@ -20,7 +20,11 @@ Internal.get_device_type(::MtlArray) = MetalDevice
 
 # unsafe_free!
 function Internal.unsafe_free_internal!(::Type{MetalDevice}, x::AbstractArray)
-    Metal.unsafe_free!(x)
+    if applicable(Metal.unsafe_free!, x)
+        Metal.unsafe_free!(x)
+    else
+        @warn "Metal.unsafe_free! is not defined for $(typeof(x))." maxlog=1
+    end
     return
 end
 

diff --git a/lib/MLDataDevices/ext/MLDataDevicesOneHotArraysExt.jl b/lib/MLDataDevices/ext/MLDataDevicesOneHotArraysExt.jl
@@ -1,13 +1,20 @@
 module MLDataDevicesOneHotArraysExt
 
 using Adapt: Adapt
-using MLDataDevices: MLDataDevices, Internal, ReactantDevice, CPUDevice
+using MLDataDevices: MLDataDevices, Internal, ReactantDevice, CPUDevice, CUDADevice,
+                     AMDGPUDevice, MetalDevice, oneAPIDevice
 using OneHotArrays: OneHotArray
 
 for op in (:get_device, :get_device_type)
     @eval Internal.$(op)(x::OneHotArray) = Internal.$(op)(x.indices)
 end
 
+for dev in (CUDADevice, AMDGPUDevice, MetalDevice, oneAPIDevice)
+    @eval function Internal.unsafe_free_internal!(D::Type{$(dev)}, x::OneHotArray)
+        return Internal.unsafe_free_internal!(D, x.indices)
+    end
+end
+
 # Reactant doesn't pay very nicely with OneHotArrays at the moment
 function Adapt.adapt_structure(dev::ReactantDevice, x::OneHotArray)
     x_cpu = Adapt.adapt_structure(CPUDevice(), x)

diff --git a/lib/MLDataDevices/ext/MLDataDevicesoneAPIExt.jl b/lib/MLDataDevices/ext/MLDataDevicesoneAPIExt.jl
@@ -31,7 +31,11 @@ Internal.get_device_type(::oneArray) = oneAPIDevice
 
 # unsafe_free!
 function Internal.unsafe_free_internal!(::Type{oneAPIDevice}, x::AbstractArray)
-    oneAPI.unsafe_free!(x)
+    if applicable(oneAPI.unsafe_free!, x)
+        oneAPI.unsafe_free!(x)
+    else
+        @warn "oneAPI.unsafe_free! is not defined for $(typeof(x))." maxlog=1
+    end
     return
 end