feat: return sequence properly + checkpointing + mincut (#1561)

* feat: return sequence properly + checkpointing + mincut

* feat: iterate over the full sequence

* test: mincut

Author: avik-pal

ext/LuxReactantExt/LuxReactantExt.jl

@@ -21,6 +21,7 @@ using Static: True, False
 
 using Lux: Lux, LuxOps, Training, Utils, StatefulLuxLayer
 using Lux.Training: TrainingBackendCache, ReactantBackend
+using Lux: get_time_dimension, time_dimension_size, init_recurrent_state
 using LuxCore: LuxCore, AbstractLuxLayer
 using MLDataDevices: MLDataDevices, ReactantDevice, get_device
 

ext/LuxReactantExt/layers.jl

@@ -1,17 +1,21 @@
 # Recurrent Layers
 function (r::Lux.Recurrence)(x::AnyTracedRArray, ps, st::NamedTuple)
     idxs = ntuple(Returns(Colon()), ndims(x) - 1)
-    N = Lux.time_dimension_size(x, r.ordering)
+    N = time_dimension_size(x, r.ordering)
 
-    (out, carry), st = r.cell(Lux.get_time_dimension(x, 1, r.ordering), ps, st)
-    sequence = similar(x, size(out)..., N)
+    # execute the first step to get the types
+    tmp = get_time_dimension(x, 1, r.ordering)
+    carry, _ = init_recurrent_state(r.cell, tmp, ps, st)
+    (tmp_result, _), _ = r.cell(tmp, ps, st)
 
-    sequence[idxs..., 1] = out
-    @trace for i in 2:N
-        (out, carry), st = r.cell((Lux.get_time_dimension(x, i, r.ordering), carry), ps, st)
+    final_result = similar(tmp_result)
+    sequence = similar(tmp_result, size(tmp_result)..., N)
+    @trace checkpointing = r.checkpointing mincut = r.mincut for i in 1:N
+        (out, carry), st = r.cell((get_time_dimension(x, i, r.ordering), carry), ps, st)
+        final_result[idxs...] = out
         sequence[idxs..., i] = out
     end
 
-    r.return_sequence isa False && return (out, st)
-    return LuxOps.eachslice(sequence, Val(ndims(sequence))), st
+    r.return_sequence isa False && return (final_result, st)
+    return eachslice(sequence; dims=ndims(sequence)), st
 end

src/layers/recurrent.jl

@@ -55,9 +55,12 @@ function init_rnn_bias(rng::AbstractRNG, init_bias, hidden_dims, bias_len)
 end
 
 """
-    Recurrence(cell;
+    Recurrence(
+        cell;
         ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex(),
-        return_sequence::Bool=false)
+        return_sequence::Bool=false,
+        mincut::Bool=false,
+    )
 
 Wraps a recurrent cell (like [`RNNCell`](@ref), [`LSTMCell`](@ref), [`GRUCell`](@ref)) to
 automatically operate over a sequence of inputs.
@@ -79,6 +82,8 @@ automatically operate over a sequence of inputs.
     the last output. Defaults to `false`.
   - `ordering`: The ordering of the batch and time dimensions in the input. Defaults to
     `BatchLastIndex()`. Alternatively can be set to `TimeLastIndex()`.
+  - `mincut`: If `true`, we will using mincut for the reverse mode differentiation.
+    *(Only for Reactant)*
 
 # Extended Help
 
@@ -119,24 +124,36 @@ struct Recurrence{R<:StaticBool,C,O<:AbstractTimeSeriesDataBatchOrdering} <:
     cell::C
     ordering::O
     return_sequence::R
+    # FIXME: checkpointing is intentionally not documented.
+    #        See https://github.com/LuxDL/Lux.jl/pull/1561#issuecomment-3564283063
+    checkpointing::Bool
+    mincut::Bool
 
     function Recurrence(
-        cell::C, ordering::AbstractTimeSeriesDataBatchOrdering, return_sequence::R
+        cell::C,
+        ordering::AbstractTimeSeriesDataBatchOrdering,
+        return_sequence::R,
+        checkpointing::Bool,
+        mincut::Bool,
     ) where {C,R}
         @assert cell isa Union{
             <:AbstractRecurrentCell,
             <:Experimental.DebugLayer{<:Any,<:Any,<:AbstractRecurrentCell},
         }
-        return new{R,C,typeof(ordering)}(cell, ordering, return_sequence)
+        return new{R,C,typeof(ordering)}(
+            cell, ordering, return_sequence, checkpointing, mincut
+        )
     end
 end
 
 function Recurrence(
     cell;
     ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex(),
     return_sequence::Bool=false,
+    checkpointing::Bool=false,
+    mincut::Bool=false,
 )
-    return Recurrence(cell, ordering, static(return_sequence))
+    return Recurrence(cell, ordering, static(return_sequence), checkpointing, mincut)
 end
 
 function (r::Recurrence)(x::AbstractArray, ps, st::NamedTuple)
@@ -233,6 +250,9 @@ function applyrecurrentcell(l::AbstractRecurrentCell, x, ps, st, carry)
 end
 applyrecurrentcell(l::AbstractRecurrentCell, x, ps, st, ::Nothing) = apply(l, x, ps, st)
 
+# Used to construct the initial state of the recurrent cell
+function init_recurrent_state end
+
 @doc doc"""
     RNNCell(in_dims => out_dims, activation=tanh; use_bias=True(), train_state=False(),
         init_bias=nothing, init_weight=nothing, init_recurrent_weight=init_weight,
@@ -343,15 +363,20 @@ end
 
 initialstates(rng::AbstractRNG, ::RNNCell) = (rng=Utils.sample_replicate(rng),)
 
-function (rnn::RNNCell{False})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(rnn::RNNCell{False}, x::AbstractMatrix, ps, st::NamedTuple)
     rng = replicate(st.rng)
     hidden_state = init_rnn_hidden_state(rng, rnn, x)
-    return rnn((x, (hidden_state,)), ps, merge(st, (; rng)))
+    return (hidden_state,), merge(st, (; rng))
 end
 
-function (rnn::RNNCell{True})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(::RNNCell{True}, x::AbstractMatrix, ps, st::NamedTuple)
     hidden_state = init_trainable_rnn_hidden_state(ps.hidden_state, x)
-    return rnn((x, (hidden_state,)), ps, st)
+    return (hidden_state,), st
+end
+
+function (rnn::RNNCell)(x::AbstractMatrix, ps, st::NamedTuple)
+    hidden_state, st = init_recurrent_state(rnn, x, ps, st)
+    return rnn((x, hidden_state), ps, st)
 end
 
 @trace function (rnn::RNNCell)(
@@ -547,31 +572,42 @@ end
 
 initialstates(rng::AbstractRNG, ::LSTMCell) = (rng=Utils.sample_replicate(rng),)
 
-function (lstm::LSTMCell{False,False})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(
+    lstm::LSTMCell{False,False}, x::AbstractMatrix, ps, st::NamedTuple
+)
     rng = replicate(st.rng)
     hidden_state = init_rnn_hidden_state(rng, lstm, x)
     memory = init_rnn_hidden_state(rng, lstm, x)
-    return lstm((x, (hidden_state, memory)), ps, merge(st, (; rng)))
+    return (hidden_state, memory), merge(st, (; rng))
 end
 
-function (lstm::LSTMCell{True,False})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(
+    lstm::LSTMCell{True,False}, x::AbstractMatrix, ps, st::NamedTuple
+)
     rng = replicate(st.rng)
     hidden_state = init_trainable_rnn_hidden_state(ps.hidden_state, x)
     memory = init_rnn_hidden_state(rng, lstm, x)
-    return lstm((x, (hidden_state, memory)), ps, merge(st, (; rng)))
+    return (hidden_state, memory), merge(st, (; rng))
 end
 
-function (lstm::LSTMCell{False,True})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(
+    lstm::LSTMCell{False,True}, x::AbstractMatrix, ps, st::NamedTuple
+)
     rng = replicate(st.rng)
     hidden_state = init_rnn_hidden_state(rng, lstm, x)
     memory = init_trainable_rnn_hidden_state(ps.memory, x)
-    return lstm((x, (hidden_state, memory)), ps, merge(st, (; rng)))
+    return (hidden_state, memory), merge(st, (; rng))
 end
 
-function (lstm::LSTMCell{True,True})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(::LSTMCell{True,True}, x::AbstractMatrix, ps, st::NamedTuple)
     hidden_state = init_trainable_rnn_hidden_state(ps.hidden_state, x)
     memory = init_trainable_rnn_hidden_state(ps.memory, x)
-    return lstm((x, (hidden_state, memory)), ps, st)
+    return (hidden_state, memory), st
+end
+
+function (lstm::LSTMCell)(x::AbstractMatrix, ps, st::NamedTuple)
+    hidden_state, st = init_recurrent_state(lstm, x, ps, st)
+    return lstm((x, hidden_state), ps, st)
 end
 
 const _LSTMCellInputType = Tuple{<:AbstractMatrix,Tuple{<:AbstractMatrix,<:AbstractMatrix}}
@@ -744,16 +780,20 @@ end
 
 initialstates(rng::AbstractRNG, ::GRUCell) = (rng=Utils.sample_replicate(rng),)
 
-function (gru::GRUCell{True})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(::GRUCell{True}, x::AbstractMatrix, ps, st::NamedTuple)
     hidden_state = init_trainable_rnn_hidden_state(ps.hidden_state, x)
-    return gru((x, (hidden_state,)), ps, st)
+    return (hidden_state,), st
 end
 
-function (gru::GRUCell{False})(x::AbstractMatrix, ps, st::NamedTuple)
+function init_recurrent_state(gru::GRUCell{False}, x::AbstractMatrix, ps, st::NamedTuple)
     rng = replicate(st.rng)
-    st = merge(st, (; rng))
     hidden_state = init_rnn_hidden_state(rng, gru, x)
-    return gru((x, (hidden_state,)), ps, st)
+    return (hidden_state,), merge(st, (; rng))
+end
+
+function (gru::GRUCell)(x::AbstractMatrix, ps, st::NamedTuple)
+    hidden_state, st = init_recurrent_state(gru, x, ps, st)
+    return gru((x, hidden_state), ps, st)
 end
 
 const _GRUCellInputType = Tuple{<:AbstractMatrix,Tuple{<:AbstractMatrix}}

test/reactant/layer_tests.jl

@@ -47,9 +47,12 @@
 
                 @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 = 1.0e-2 rtol = 1.0e-2
-                    @test check_approx(∂ps_ra, ∂ps; atol=1.0e-2, rtol=1.0e-2)
+                    @testset for mincut in (true, false), checkpointing in (false,)
+                        model_ = Recurrence(cell(4 => 4); ordering, mincut, checkpointing)
+                        ∂x_ra, ∂ps_ra = @jit ∇sumabs2_enzyme(model_, x_ra, ps_ra, st_ra)
+                        @test ∂x_ra ≈ ∂x atol = 1.0e-2 rtol = 1.0e-2
+                        @test check_approx(∂ps_ra, ∂ps; atol=1.0e-2, rtol=1.0e-2)
+                    end
                 end
 
                 model2 = Recurrence(cell(4 => 4); ordering, return_sequence=true)