feat: don't unroll Recurrence

Project.toml

@@ -1,7 +1,7 @@
 name = "Lux"
 uuid = "b2108857-7c20-44ae-9111-449ecde12c47"
 authors = ["Avik Pal <[email protected]> and contributors"]
-version = "1.6.0"
+version = "1.7.0"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
@@ -100,7 +100,7 @@ LinearAlgebra = "1.10"
 LossFunctions = "0.11.1, 1"
 LuxCore = "1.2"
 LuxLib = "1.5.0"
-MLDataDevices = "1.6.6"
+MLDataDevices = "1.6.10"
 MLUtils = "0.4.4"
 MPI = "0.20.19"
 MacroTools = "0.5.13"
@@ -110,7 +110,7 @@ NNlib = "0.9.26"
 Optimisers = "0.4.1"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Reexport = "1.2.2"
 ReverseDiff = "1.15"
 SIMDTypes = "0.1"

docs/Project.toml

@@ -50,13 +50,13 @@ LuxCUDA = "0.3.2"
 LuxCore = "1.2"
 LuxLib = "1.3.4"
 LuxTestUtils = "1.5"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 NNlib = "0.9.26"
 Optimisers = "0.4.1"
 Pkg = "1.10"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 StableRNGs = "1"
 StaticArrays = "1"
 WeightInitializers = "1"

examples/CIFAR10/Project.toml

@@ -39,7 +39,7 @@ Optimisers = "0.4.1"
 Printf = "1.10"
 ProgressTables = "0.1"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 StableRNGs = "1.0.2"
 Statistics = "1.10"
 Zygote = "0.6.70, 0.7"

examples/ConvolutionalVAE/Project.toml

@@ -26,4 +26,4 @@ OneHotArrays = "0.2.6"
 Optimisers = "0.4"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"

examples/GCN_Cora/Project.toml

@@ -21,5 +21,5 @@ OneHotArrays = "0.2"
 Optimisers = "0.4.4"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Statistics = "1.10"

examples/ImageNet/Project.toml

@@ -33,7 +33,7 @@ ImageMagick = "1"
 JLD2 = "0.5.1"
 Lux = "1"
 LuxCUDA = "0.3.3"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 MLUtils = "0.4.4"
 MPI = "0.20.21"
 NCCL = "0.1.1"

examples/PINN2DPDE/Project.toml

@@ -21,5 +21,5 @@ OnlineStats = "1.7.1"
 Optimisers = "0.4.1"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Statistics = "1.10"

examples/PolynomialFitting/Project.toml

@@ -15,5 +15,5 @@ Lux = "1"
 Optimisers = "0.4.1"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Statistics = "1.10"

examples/RealNVP/Project.toml

@@ -19,5 +19,5 @@ MLUtils = "0.4.5"
 Optimisers = "0.4.4"
 Printf = "1.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Statistics = "1.10"

examples/SimpleChains/Project.toml

@@ -19,6 +19,6 @@ MLUtils = "0.4"
 OneHotArrays = "0.2.5"
 Optimisers = "0.4.1"
 Random = "1"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 SimpleChains = "0.4.6"
 Zygote = "0.6.69, 0.7"

examples/SimpleRNN/Project.toml

@@ -14,4 +14,4 @@ JLD2 = "0.5"
 Lux = "1"
 MLUtils = "0.4"
 Optimisers = "0.4.1"
-Reactant = "0.2.21"
+Reactant = "0.2.27"

ext/LuxReactantExt/LuxReactantExt.jl

@@ -2,12 +2,14 @@ module LuxReactantExt
 
 using Enzyme: Enzyme, Const, Duplicated, Active
 using Optimisers: Optimisers
-using Reactant: Reactant, @compile, @code_hlo, AnyTracedRArray, TracedRArray, TracedRNumber
+using Reactant: Reactant, @compile, @code_hlo, @trace, AnyTracedRArray, TracedRArray,
+                TracedRNumber
 using Setfield: @set!
 using Static: True, False
 
 using Lux: Lux, LuxOps, Training, Utils
 using Lux.Training: TrainingBackendCache, ReactantBackend
+using LuxCore: LuxCore
 
 Lux.is_extension_loaded(::Val{:Reactant}) = true
 
@@ -26,5 +28,6 @@ end
 
 include("patches.jl")
 include("training.jl")
+include("layers.jl")
 
 end

ext/LuxReactantExt/layers.jl

@@ -0,0 +1,54 @@
+# Embedding
+function (e::Lux.Embedding)(x::TracedRNumber{<:Reactant.ReactantInt}, ps, st::NamedTuple)
+    return ps.weight[:, x], st
+end
+
+# Recurrent Layers
+# TODO: Once we can eliminate dead-args in while loop we should remove this case and only
+#       use the later function for maintenance purposes.
+function (r::Lux.Recurrence{False})(x::AnyTracedRArray, ps, st::NamedTuple)
+    if r.ordering isa Lux.TimeLastIndex ||
+       (r.ordering isa Lux.BatchLastIndex && ndims(x) == 2)
+        idxs = ntuple(Returns(Colon()), ndims(x) - 1)
+        (out, carry), st = r.cell(x[idxs..., 1], ps, st)
+        @trace for i in 2:size(x, ndims(x))
+            (out, carry), st = r.cell((x[idxs..., i], carry), ps, st)
+        end
+        return out, st
+    elseif r.ordering isa Lux.BatchLastIndex
+        idxs = ntuple(Returns(Colon()), ndims(x) - 2)
+        (out, carry), st = r.cell(x[idxs..., 1, :], ps, st)
+        @trace for i in 2:size(x, ndims(x) - 1)
+            (out, carry), st = r.cell((x[idxs..., i, :], carry), ps, st)
+        end
+        return out, st
+    else
+        error("Unknown ordering: $(r.ordering)")
+    end
+end
+
+function (r::Lux.Recurrence{True})(x::AnyTracedRArray, ps, st::NamedTuple)
+    if r.ordering isa Lux.TimeLastIndex ||
+       (r.ordering isa Lux.BatchLastIndex && ndims(x) == 2)
+        idxs = ntuple(Returns(Colon()), ndims(x) - 1)
+        (out, carry), st = r.cell(x[idxs..., 1], ps, st)
+        sequence = similar(out, size(out)..., size(x, ndims(x)))
+        sequence[idxs..., 1] .= out
+        @trace for i in 2:size(x, ndims(x))
+            (out, carry), st = r.cell((x[idxs..., i], carry), ps, st)
+            sequence[idxs..., i] = out
+        end
+    elseif r.ordering isa Lux.BatchLastIndex
+        idxs = ntuple(Returns(Colon()), ndims(x) - 2)
+        (out, carry), st = r.cell(x[idxs..., 1, :], ps, st)
+        sequence = similar(out, size(out)..., size(x, ndims(x) - 1))
+        sequence[idxs..., :, 1] .= out
+        @trace for i in 2:size(x, ndims(x) - 1)
+            (out, carry), st = r.cell((x[idxs..., i, :], carry), ps, st)
+            sequence[idxs..., :, i] = out
+        end
+    else
+        error("Unknown ordering: $(r.ordering)")
+    end
+    return (out, eachslice(sequence; dims=ndims(sequence))), st
+end

ext/LuxReactantExt/patches.jl

@@ -2,8 +2,3 @@ Utils.vec(x::AnyTracedRArray) = Reactant.TracedUtils.materialize_traced_array(ve
 
 # XXX: Use PoolDims once EnzymeJAX supports stablehlo.reduce_window adjoint
 Lux.calculate_pool_dims(g::Lux.GlobalPoolMode, ::TracedRArray) = g
-
-# Embedding
-function (e::Lux.Embedding)(x::TracedRNumber{<:Reactant.ReactantInt}, ps, st::NamedTuple)
-    return ps.weight[:, x], st
-end

lib/LuxCore/Project.toml

@@ -38,9 +38,9 @@ Compat = "4.16"
 DispatchDoctor = "0.4.10"
 EnzymeCore = "0.8.6"
 Functors = "0.5"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 ReverseDiff = "1.15"
 Setfield = "1"
 Tracker = "0.2.36"

lib/LuxCore/test/Project.toml

@@ -14,7 +14,7 @@ Aqua = "0.8.7"
 EnzymeCore = "0.8.6"
 ExplicitImports = "1.9.0"
 Functors = "0.5"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 Optimisers = "0.3.4, 0.4"
 Random = "1.10"
 Test = "1.10"

lib/LuxLib/Project.toml

@@ -85,14 +85,14 @@ LinearAlgebra = "1.10"
 LoopVectorization = "0.12.171"
 LuxCore = "1.2"
 MKL = "0.7, 0.8"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 Markdown = "1.10"
 NNlib = "0.9.26"
 Octavian = "0.3.28"
 Polyester = "0.7.15"
 Preferences = "1.4.3"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 Reexport = "1"
 ReverseDiff = "1.15"
 SLEEFPirates = "0.6.43"

lib/LuxLib/test/Project.toml

@@ -48,7 +48,7 @@ JLArrays = "0.1.5"
 LoopVectorization = "0.12.171"
 LuxTestUtils = "1.5"
 MKL = "0.7"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 NNlib = "0.9.26"
 Octavian = "0.3.28"
 Pkg = "1.10"

lib/LuxTestUtils/Project.toml

@@ -35,7 +35,7 @@ FiniteDiff = "2.23.1"
 ForwardDiff = "0.10.36"
 Functors = "0.5"
 JET = "0.9.6"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 Optimisers = "0.3.4, 0.4"
 ReverseDiff = "1.15.3"
 Test = "1.10"

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.9"
+version = "1.6.10"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -66,7 +66,7 @@ Metal = "1"
 OneHotArrays = "0.2.5"
 Preferences = "1.4"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 RecursiveArrayTools = "3.8"
 ReverseDiff = "1.15"
 SparseArrays = "1.10"

lib/MLDataDevices/ext/MLDataDevicesReactantExt.jl

@@ -16,7 +16,7 @@ end
 
 # Query Device from Array
 function Internal.get_device(x::Union{ConcreteRNumber, ConcreteRArray})
-    return ReactantDevice(XLA.client(x.data), XLA.device(x.data))
+    return ReactantDevice(XLA.client(x), XLA.device(x))
 end
 
 function Internal.get_device(::Union{TracedRArray, TracedRNumber})

lib/WeightInitializers/Project.toml

@@ -41,7 +41,7 @@ GPUArraysCore = "0.1.6, 0.2"
 LinearAlgebra = "1.10"
 Metal = "1.3.0"
 Random = "1.10"
-Reactant = "0.2.21"
+Reactant = "0.2.27"
 SpecialFunctions = "2.4"
 Statistics = "1.10"
 julia = "1.10"

test/Project.toml

@@ -60,7 +60,7 @@ LoopVectorization = "0.12.171"
 LuxCore = "1.2"
 LuxLib = "1.3.4"
 LuxTestUtils = "1.5"
-MLDataDevices = "1.6"
+MLDataDevices = "1.6.10"
 MLUtils = "0.4.3"
 NNlib = "0.9.26"
 Octavian = "0.3.28"

test/reactant/layer_tests.jl

@@ -28,7 +28,8 @@ end
     using Reactant, Lux
     using LuxTestUtils: check_approx
 
-    rng = StableRNG(123)
+    # StableRNG uses UInt128 for seed that is not supported by Reactant inside loops
+    rng = Random.default_rng()
 
     @testset "$(mode)" for (mode, atype, dev, ongpu) in MODES
         if mode == "amdgpu"
@@ -45,22 +46,28 @@ end
         dev = reactant_device(; force=true)
 
         @testset for cell in (RNNCell, LSTMCell, GRUCell)
-            model = Recurrence(cell(4 => 4))
-            ps, st = Lux.setup(rng, model)
-            ps_ra, st_ra = (ps, st) |> dev
-            x = rand(Float32, 4, 16, 12)
-            x_ra = x |> dev
-
-            y_ra, _ = @jit model(x_ra, ps_ra, st_ra)
-            y, _ = model(x, ps, st)
-
-            @test y_ra≈y atol=1e-2 rtol=1e-2
-
-            @testset "gradient" begin
-                ∂x, ∂ps = ∇sumabs2_zygote(model, x, ps, st)
-                ∂x_ra, ∂ps_ra = @jit ∇sumabs2_enzyme(model, x_ra, ps_ra, st_ra)
-                @test ∂x_ra≈∂x atol=1e-2 rtol=1e-2
-                @test check_approx(∂ps_ra, ∂ps; atol=1e-2, rtol=1e-2)
+            @testset for ordering in (BatchLastIndex(), TimeLastIndex())
+                model = Recurrence(cell(4 => 4); ordering)
+                ps, st = Lux.setup(rng, model)
+                ps_ra, st_ra = (ps, st) |> dev
+                if ordering isa TimeLastIndex
+                    x = rand(Float32, 4, 12, 16)
+                else
+                    x = rand(Float32, 4, 16, 12)
+                end
+                x_ra = x |> dev
+
+                y_ra, _ = @jit model(x_ra, ps_ra, st_ra)
+                y, _ = model(x, ps, st)
+
+                @test y_ra≈y atol=1e-2 rtol=1e-2
+
+                @testset "gradient" begin
+                    ∂x, ∂ps = ∇sumabs2_zygote(model, x, ps, st)
+                    ∂x_ra, ∂ps_ra = @jit ∇sumabs2_enzyme(model, x_ra, ps_ra, st_ra)
+                    @test ∂x_ra≈∂x atol=1e-2 rtol=1e-2
+                    @test check_approx(∂ps_ra, ∂ps; atol=1e-2, rtol=1e-2)
+                end
             end
         end
     end