Recurse into parent

LuxDL · Mar 16, 2024 · 8fa28e9 · 8fa28e9
1 parent de25a7a
commit 8fa28e9
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxDeviceUtils"
 uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.1.17"
+version = "0.1.18"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/src/LuxDeviceUtils.jl b/src/LuxDeviceUtils.jl
@@ -318,7 +318,16 @@ end
 Returns the device of the array `x`. Trigger Packages must be loaded for this to return the
 correct device.
 """
-get_device(::AbstractArray) = LuxCPUDevice()
+function get_device(x::AbstractArray)
+    if hasmethod(parent, Tuple{typeof(x)})
+        parent_x = parent(x)
+        parent_x === x && return LuxCPUDevice()
+        return get_device(parent_x)
+    end
+    return LuxCPUDevice()
+end
+
+CRC.@non_differentiable get_device(::Any...)
 
 # Adapt Interface
 abstract type AbstractLuxDeviceAdaptor end