feat: add warning on attempting to move architecture

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "0.1.21"
+version = "0.1.22"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
@@ -12,10 +12,12 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 
 [extensions]
 LuxCoreChainRulesCoreExt = "ChainRulesCore"
+LuxCoreMLDataDevicesExt = "MLDataDevices"
 LuxCoreEnzymeCoreExt = "EnzymeCore"
 
 [compat]
@@ -26,6 +28,7 @@ DispatchDoctor = "0.4.10"
 EnzymeCore = "0.7.7"
 ExplicitImports = "1.9.0"
 Functors = "0.4.8"
+MLDataDevices = "1"
 Optimisers = "0.3"
 Random = "1.10"
 Setfield = "1"
@@ -36,9 +39,10 @@ julia = "1.10"
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7"
+MLDataDevices = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 Optimisers = "3bd65402-5787-11e9-1adc-39752487f4e2"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "EnzymeCore", "ExplicitImports", "Optimisers", "Random", "Test"]
+test = ["Aqua", "EnzymeCore", "ExplicitImports", "MLDataDevices", "Optimisers", "Random", "Test"]

ext/LuxCoreMLDataDevicesExt.jl

@@ -0,0 +1,16 @@
+module LuxCoreMLDataDevicesExt
+
+using LuxCore: LuxCore
+using MLDataDevices: MLDataDevices
+
+for (dev) in (:CPU, :CUDA, :AMDGPU, :Metal, :oneAPI)
+    ldev = Symbol(dev, :Device)
+    @eval function (::MLDataDevices.$(ldev))(NN::LuxCore.AbstractExplicitLayer)
+        @warn "Lux layers are stateless and hence don't participate in device transfers. \
+               Apply this function on the parameters and states generated using \
+               `LuxCore.setup`."
+        return NN
+    end
+end
+
+end

test/runtests.jl

@@ -1,4 +1,5 @@
-using Aqua, ExplicitImports, Functors, LuxCore, Optimisers, Random, Test, EnzymeCore
+using Aqua, ExplicitImports, Functors, LuxCore, Optimisers, Random, Test, EnzymeCore,
+      MLDataDevices
 
 rng = LuxCore._default_rng()
 
@@ -290,4 +291,14 @@ end
         @test_throws ArgumentError BatchDuplicatedNoNeed(d, (d, d))
         @test Const(d) isa Const
     end
+
+    @testset "Device Transfer Warnings" begin
+        my_layer = Dense(2, 2)
+
+        dev = cpu_device()
+        @test_logs (
+            :warn, "Lux layers are stateless and hence don't participate in device \
+                    transfers. Apply this function on the parameters and states generated \
+                    using `LuxCore.setup`.") dev(my_layer)
+    end
 end