fix: update tracing to new Reactant API (#1636)

Author: avik-pal

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.29.3"
+version = "1.29.4"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -89,7 +89,7 @@ ZygoteExt = "Zygote"
 
 [compat]
 ADTypes = "1.15"
-Adapt = "4.1"
+Adapt = "4.4"
 ArrayInterface = "7.17.1"
 CUDA = "5.8"
 ChainRulesCore = "1.25.1"
@@ -122,12 +122,12 @@ Optimisers = "0.4.6"
 PrecompileTools = "1.2.1"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.190"
+Reactant = "0.2.200"
 ReactantCore = "0.1.16"
 Reexport = "1.2.2"
 ReverseDiff = "1.15"
 SciMLPublic = "1.0.0"
-Setfield = "1.1.1"
+Setfield = "1.1.2"
 SimpleChains = "0.4.7"
 Static = "1.1.1"
 StaticArraysCore = "1.4.3"

docs/Project.toml

@@ -68,7 +68,7 @@ NPZ = "0.4.3"
 Optimisers = "0.4.6"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.173"
+Reactant = "0.2.200"
 StableRNGs = "1"
 StaticArrays = "1"
 WeightInitializers = "1.3"

ext/ReactantExt/tracing.jl

@@ -3,12 +3,14 @@ function Reactant.traced_type_inner(
     seen,
     @nospecialize(mode::Reactant.TraceMode),
     @nospecialize(track_numbers::Type),
-    @nospecialize(sharding),
+    @nospecialize(ndevices),
     @nospecialize(runtime)
 ) where {ST,M,psT,stT}
-    args = (Val(mode), track_numbers, sharding, runtime)
     return StatefulLuxLayer{
-        ST,M,Reactant.traced_type(psT, args...),Reactant.traced_type(stT, args...)
+        ST,
+        M,
+        Reactant.traced_type_inner(psT, seen, mode, track_numbers, ndevices, runtime),
+        Reactant.traced_type_inner(stT, seen, mode, track_numbers, ndevices, runtime),
     }
 end
 

lib/LuxCore/Project.toml

@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.5.2"
+version = "1.5.3"
 
 [deps]
 DispatchDoctor = "8d63f2c5-f18a-4cf2-ba9d-b3f60fc568c8"
@@ -33,7 +33,7 @@ ReactantExt = "Reactant"
 SetfieldExt = "Setfield"
 
 [compat]
-Adapt = "4.1"
+Adapt = "4.4"
 ArrayInterface = "7.9"
 ChainRulesCore = "1.25.1"
 DispatchDoctor = "0.4.10"
@@ -42,7 +42,7 @@ Flux = "0.16.3"
 Functors = "0.5"
 MLDataDevices = "1.17"
 Random = "1.10"
-Reactant = "0.2.170"
+Reactant = "0.2.200"
 ReverseDiff = "1.15"
 SciMLPublic = "1.0.0"
 Setfield = "1"

lib/LuxCore/ext/ReactantExt.jl

@@ -15,7 +15,7 @@ function Reactant.traced_type_inner(
     seen,
     @nospecialize(mode::Reactant.TraceMode),
     @nospecialize(track_numbers::Type),
-    @nospecialize(sharding),
+    @nospecialize(ndevices),
     @nospecialize(runtime)
 )
     return T

lib/LuxTestUtils/Project.toml

@@ -26,7 +26,7 @@ MLDataDevices = {path = "../MLDataDevices"}
 
 [compat]
 ADTypes = "1.10"
-Adapt = "4.1"
+Adapt = "4.4"
 ArrayInterface = "7.17.1"
 ChainRulesCore = "1.25.1"
 ComponentArrays = "0.15.22"

lib/MLDataDevices/Project.toml

@@ -1,7 +1,7 @@
 name = "MLDataDevices"
 uuid = "7e8f7934-dd98-4c1a-8fe8-92b47a384d40"
 authors = ["Avik Pal <avikpal@mit.edu> and contributors"]
-version = "1.17.1"
+version = "1.17.2"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -54,7 +54,7 @@ OpenCLExt = ["GPUArrays", "OpenCL"]
 
 [compat]
 AMDGPU = "1, 2"
-Adapt = "4.1"
+Adapt = "4.4"
 CUDA = "5.8"
 ChainRules = "1.51"
 ChainRulesCore = "1.25.1"
@@ -67,7 +67,7 @@ Metal = "1"
 OneHotArrays = "0.2.10"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.170"
+Reactant = "0.2.200"
 RecursiveArrayTools = "3.8"
 ReverseDiff = "1.15"
 SciMLPublic = "1.0.0"

lib/MLDataDevices/ext/ReactantExt.jl

@@ -15,7 +15,7 @@ function Reactant.traced_type_inner(
     seen,
     @nospecialize(mode::Reactant.TraceMode),
     @nospecialize(track_numbers::Type),
-    @nospecialize(sharding),
+    @nospecialize(ndevices),
     @nospecialize(runtime)
 )
     return T