Recurse into parent

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

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -182,6 +182,6 @@ steps:
               - "1"
 
 env:
-  RETESTITEMS_NWORKERS: 4
+  RETESTITEMS_NWORKERS: 2
   RETESTITEMS_NWORKER_THREADS: 2
   SECRET_CODECOV_TOKEN: "PxSr3Y7vdbiwaoX51uGykPsogxmP1IOBt5Z8TwP9GqDxIrvFocEVV2DR4Bebee12G/HYvXtQTyYXH49DpzlsfJ7ri1GQZxd9WRr+aM1DDYmzfDCfpadp4hMoJ5NQvmc/PzeGrNWOOaewaLTUP1eEaG4suygZN0lc5q9BCchIJeqoklGms5DVt/HtfTmwoD/s4wGoIJINi4RoFgnCAkzSh11hTAkyjVerfBGWEi/8E6+WBq3UKwaW4HnT02wG9qFnD4XkHpIpjMxJTpdBn5ufKI+QoJ7qJHlwqgDCtsOCblApccLTjH/BnTahNoSb/b0wdS/cblOTrtdPGzZ5UvmQ4Q==;U2FsdGVkX1/Ji2Nqeq3tqTYCBik6iXILP+rriPRqj/qxhFu4vBWWT3UnlfqDzj6oVdXyuKt0+5e+x33x2S0mBw=="
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