fix: use functors for testing wrapped arrays

LuxDL · Dec 14, 2024 · c462b7d · c462b7d
1 parent fdb0170
commit c462b7d
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 3 deletions.
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.6.4"
+version = "1.6.5"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

diff --git a/lib/MLDataDevices/src/public.jl b/lib/MLDataDevices/src/public.jl
@@ -399,5 +399,7 @@ If `MLDataDevices.isleaf(x::T)` is not defined, then it will fall back to `Funct
 """
 isleaf(x) = Functors.isleaf(x)
 
-isleaf(::AbstractArray{T}) where {T} = isbitstype(T) || T <: Number # BigFloat and such are not bitstype
-isleaf(::Adapt.WrappedArray) = false
+function isleaf(x::AbstractArray{T}) where {T}
+    parent(x) !== x && return Functors.isleaf(x)
+    return isbitstype(T) || T <: Number # BigFloat and such are not bitstype
+end