feat: tracing stateful layers (#1257)

* feat: tracing stateful layers * test: StatefulLuxLayer tracing
LuxDL · Mar 7, 2025 · 16031b6 · 16031b6 · avik-pal · Mar 7, 2025
1 parent 3c93c68
commit 16031b6
Show file tree

Hide file tree

Showing 5 changed files with 92 additions and 2 deletions.
diff --git a/ext/LuxReactantExt/LuxReactantExt.jl b/ext/LuxReactantExt/LuxReactantExt.jl
@@ -7,7 +7,7 @@ using Reactant: Reactant, @compile, @code_hlo, @trace, AnyTracedRArray, TracedRA
 using Setfield: @set!
 using Static: True, False
 
-using Lux: Lux, LuxOps, Training, Utils
+using Lux: Lux, LuxOps, Training, Utils, StatefulLuxLayer
 using Lux.Training: TrainingBackendCache, ReactantBackend
 using LuxCore: LuxCore
 
@@ -29,5 +29,6 @@ end
 include("patches.jl")
 include("training.jl")
 include("layers.jl")
+include("tracing.jl")
 
 end
diff --git a/ext/LuxReactantExt/tracing.jl b/ext/LuxReactantExt/tracing.jl
@@ -0,0 +1,48 @@
+@static if isdefined(Reactant, :ConcreteIFRTArray)
+    function Reactant.traced_type_inner(
+            @nospecialize(T::Type{StatefulLuxLayer{ST, M, psT, stT}}),
+            seen,
+            @nospecialize(mode::Reactant.TraceMode),
+            @nospecialize(track_numbers::Type),
+            @nospecialize(sharding),
+            @nospecialize(runtime)
+        ) where {ST, M, psT, stT}
+        return StatefulLuxLayer{
+            ST,
+            M,
+            Reactant.traced_type(psT. seen, mode, track_numbers, sharding, runtime),
+            Reactant.traced_type(stT, seen, mode, track_numbers, sharding, runtime),
+        }
+    end
+else
+    function Reactant.traced_type_inner(
+            @nospecialize(T::Type{StatefulLuxLayer{ST, M, psT, stT}}),
+            seen,
+            @nospecialize(mode::Reactant.TraceMode),
+            @nospecialize(track_numbers::Type),
+            @nospecialize(sharding)
+        ) where {ST, M, psT, stT}
+        return StatefulLuxLayer{
+            ST,
+            M,
+            Reactant.traced_type(psT, seen, mode, track_numbers, sharding),
+            Reactant.traced_type(stT, seen, mode, track_numbers, sharding),
+        }
+    end
+end
+
+function Reactant.make_tracer(
+        seen,
+        @nospecialize(model::StatefulLuxLayer),
+        @nospecialize(path),
+        mode;
+        kwargs...
+    )
+    return StatefulLuxLayer(
+        model.model,
+        Reactant.make_tracer(seen, model.ps, (path..., :ps), mode; kwargs...),
+        Reactant.make_tracer(seen, model.st, (path..., :st), mode; kwargs...),
+        Reactant.make_tracer(seen, model.st_any, (path..., :st_any), mode; kwargs...),
+        model.fixed_state_type
+    )
+end
diff --git a/lib/LuxCore/Project.toml b/lib/LuxCore/Project.toml
@@ -1,7 +1,7 @@
 name = "LuxCore"
 uuid = "bb33d45b-7691-41d6-9220-0943567d0623"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.2.2"
+version = "1.2.3"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

diff --git a/lib/LuxCore/ext/LuxCoreReactantExt.jl b/lib/LuxCore/ext/LuxCoreReactantExt.jl
@@ -10,6 +10,25 @@ function Reactant.make_tracer(
     return model
 end
 
+@static if isdefined(Reactant, :ConcreteIFRTArray)
+    function Reactant.traced_type_inner(
+            @nospecialize(T::Type{<:AbstractLuxLayer}), seen,
+            @nospecialize(mode::Reactant.TraceMode),
+            @nospecialize(track_numbers::Type), @nospecialize(sharding),
+            @nospecialize(runtime)
+        )
+        return T
+    end
+else
+    function Reactant.traced_type_inner(
+            @nospecialize(T::Type{<:AbstractLuxLayer}), seen,
+            @nospecialize(mode::Reactant.TraceMode),
+            @nospecialize(track_numbers::Type), @nospecialize(sharding)
+        )
+        return T
+    end
+end
+
 LuxCore.replicate(rng::Reactant.TracedRNG) = copy(rng)
 LuxCore.replicate(rng::Reactant.ConcreteRNG) = copy(rng)
 

diff --git a/test/reactant/tracing_tests.jl b/test/reactant/tracing_tests.jl
@@ -0,0 +1,22 @@
+@testitem "Tracing" tags = [:reactant] skip = :(Sys.iswindows()) begin
+    using Reactant, Lux, Random
+
+    model = Chain(Dense(2 => 3, relu), BatchNorm(3), Dense(3 => 2))
+    ps, st = Lux.setup(Random.default_rng(), model)
+
+    smodel = StatefulLuxLayer{true}(model, ps, st)
+    smodel_ra = Reactant.to_rarray(smodel)
+
+    @test get_device_type(smodel_ra.ps) <: ReactantDevice
+    @test get_device_type(smodel_ra.st) <: ReactantDevice
+    @test smodel_ra.st_any === nothing
+    @test smodel_ra.fixed_state_type == smodel.fixed_state_type
+
+    smodel = StatefulLuxLayer{false}(model, ps, st)
+    smodel_ra = Reactant.to_rarray(smodel)
+
+    @test get_device_type(smodel_ra.ps) <: ReactantDevice
+    @test get_device_type(smodel_ra.st_any) <: ReactantDevice
+    @test smodel_ra.st === nothing
+    @test smodel_ra.fixed_state_type == smodel.fixed_state_type
+end