Add fast and type stable paths for certain datastructures

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxDeviceUtils"
 uuid = "34f89e08-e1d5-43b4-8944-0b49ac560553"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "0.1.7"
+version = "0.1.8"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

src/LuxDeviceUtils.jl

@@ -209,14 +209,25 @@ Return a `LuxCPUDevice` object which can be used to transfer data to CPU.
 """
 @inline cpu_device() = LuxCPUDevice()
 
-(::LuxCPUDevice)(x) = fmap(Base.Fix1(adapt, LuxCPUAdaptor()), x; exclude=_isleaf)
-(::LuxCUDADevice)(x) = fmap(Base.Fix1(adapt, LuxCUDAAdaptor()), x; exclude=_isleaf)
-(::LuxAMDGPUDevice)(x) = fmap(Base.Fix1(adapt, LuxAMDGPUAdaptor()), x; exclude=_isleaf)
-(::LuxMetalDevice)(x) = fmap(Base.Fix1(adapt, LuxMetalAdaptor()), x; exclude=_isleaf)
-
-for dev in (LuxCPUDevice, LuxCUDADevice, LuxAMDGPUDevice, LuxMetalDevice)
+# Dispatches for Different Data Structures
+# Abstract Array / Tuples / NamedTuples have special fast paths to facilitate type stability
+# For all other types we rely on fmap which means we lose type stability.
+# For Lux, typically models only has these 3 datastructures so we should be mostly fine.
+for (dev) in (:CPU, :CUDA, :AMDGPU, :Metal)
+    ldev = Symbol("Lux$(dev)Device")
+    ladaptor = Symbol("Lux$(dev)Adaptor")
     @eval begin
-        function (::$dev)(::LuxCore.AbstractExplicitLayer)
+        function (::$(ldev))(x::AbstractArray)
+            fn = Base.Fix1(adapt, $(ladaptor)())
+            return _isbitsarray(x) ? fn(x) : map(fn, x)
+        end
+        (::$(ldev))(x::Tuple) = map(Base.Fix1(adapt, $(ladaptor)()), x)
+        (::$(ldev))(x::NamedTuple{F}) where {F} = NamedTuple{F}($(ldev)(values(x)))
+        function (::$(ldev))(x)
+            _isleaf(x) && return adapt($(ladaptor)(), x)
+            return fmap(Base.Fix1(adapt, $(ladaptor)()), x; exclude=_isleaf)
+        end
+        function (::$(ldev))(::LuxCore.AbstractExplicitLayer)
             throw(ArgumentError("Lux layers are stateless and hence don't participate in device transfers. Apply this function on the parameters and states generated using `Lux.setup`."))
         end
     end