feat: bidirectional RNN + debugging RNNs (#708)

* feat: add BidirectionalRNN

* fix: avoid lazy reverse

* fix: handle reverse for operator overloading AD

* fix: allow debug modes for recurrent layers

* fix: soa/aos handling for multigate

---------

Co-authored-by: Avik Pal <[email protected]>

docs/src/api/Lux/layers.md

@@ -52,6 +52,7 @@ LSTMCell
 RNNCell
 Recurrence
 StatefulRecurrentCell
+BidirectionalRNN
 ```
 
 ## Linear Layers

ext/LuxReverseDiffExt/LuxReverseDiffExt.jl

@@ -26,6 +26,25 @@ Lux.apply(m::Lux.AbstractExplicitLayer, x::TrackedArray, ps, st) = m(x, ps, st)
 @inline Lux.__eltype(::TrackedReal{T}) where {T} = T
 @inline Lux.__eltype(::AbstractArray{<:TrackedReal{T}}) where {T} = T
 
+@inline Lux.__reverse(x::TrackedArray; dims=:) = ArrayInterface.aos_to_soa(reverse(x; dims))
+@inline function Lux.__reverse(x::AbstractArray{<:TrackedReal}; dims=:)
+    return ArrayInterface.aos_to_soa(reverse(x; dims))
+end
+
+# multigate: avoid soa formation
+@inline function Lux._gate(x::TrackedArray{T, R, 1}, h::Int, n::Int) where {T, R}
+    return x[Lux._gate(h, n)]
+end
+@inline function Lux._gate(x::AbstractVector{<:TrackedReal}, h::Int, n::Int)
+    return ArrayInterface.aos_to_soa(view(x, Lux._gate(h, n)))
+end
+@inline function Lux._gate(x::TrackedArray{T, R, 2}, h::Int, n::Int) where {T, R}
+    return x[Lux._gate(h, n), :]
+end
+@inline function Lux._gate(x::AbstractMatrix{<:TrackedReal}, h::Int, n::Int)
+    return ArrayInterface.aos_to_soa(view(x, Lux._gate(h, n), :))
+end
+
 @inline function Lux.__convert_eltype(::Type{T}, x::AbstractArray{<:TrackedReal}) where {T}
     @warn "`Lux.__convert_eltype` doesn't support converting element types of ReverseDiff \
            `TrackedReal` arrays. Currently this is a no-op." maxlog=1

ext/LuxTrackerExt.jl

@@ -65,6 +65,11 @@ Lux.apply(m::Lux.AbstractExplicitLayer, x::TrackedArray, ps, st) = m(x, ps, st)
 @inline Lux.__eltype(::TrackedReal{T}) where {T} = T
 @inline Lux.__eltype(::AbstractArray{<:TrackedReal{T}}) where {T} = T
 
+@inline Lux.__reverse(x::TrackedArray; dims=:) = ArrayInterface.aos_to_soa(reverse(x; dims))
+@inline function Lux.__reverse(x::AbstractArray{<:TrackedReal}; dims=:)
+    return ArrayInterface.aos_to_soa(reverse(x; dims))
+end
+
 # SimpleChains.jl: DON'T REPLACE THESE WITH @grad_from_chainrules
 for T1 in (:TrackedArray, :AbstractArray), T2 in (:TrackedArray, :AbstractArray)
     T1 === :AbstractArray && T2 === :AbstractArray && continue

src/Lux.jl

@@ -42,6 +42,9 @@ include("preferences.jl")
 include("custom_errors.jl")
 include("utils.jl")
 
+# Experimental
+include("contrib/contrib.jl")
+
 # Layer Implementations
 include("layers/basic.jl")
 include("layers/containers.jl")
@@ -54,9 +57,6 @@ include("layers/extension.jl")
 # Pretty Printing
 include("layers/display.jl")
 
-# Experimental
-include("contrib/contrib.jl")
-
 # Helpful Functionalities
 include("helpers/stateful.jl")
 include("helpers/compact.jl")
@@ -93,7 +93,7 @@ export AlphaDropout, Dropout, VariationalHiddenDropout
 export BatchNorm, GroupNorm, InstanceNorm, LayerNorm
 export WeightNorm
 export NoOpLayer, ReshapeLayer, SelectDim, FlattenLayer, WrappedFunction, ReverseSequence
-export RNNCell, LSTMCell, GRUCell, Recurrence, StatefulRecurrentCell
+export RNNCell, LSTMCell, GRUCell, Recurrence, StatefulRecurrentCell, BidirectionalRNN
 export SamePad, TimeLastIndex, BatchLastIndex
 
 export StatefulLuxLayer

src/contrib/debug.jl

@@ -45,7 +45,7 @@ See [`Lux.Experimental.@debug_mode`](@ref) to construct this layer.
 """
 @concrete struct DebugLayer{NaNCheck, ErrorCheck} <:
                  AbstractExplicitContainerLayer{(:layer,)}
-    layer
+    layer <: AbstractExplicitLayer
     location::KeyPath
 end
 

src/layers/basic.jl

@@ -86,22 +86,22 @@ end
 
 @inline function (r::ReverseSequence{Nothing})(
         x::AbstractVector{T}, ps, st::NamedTuple) where {T}
-    return (isbitstype(T) ? reverse(x) : Iterators.reverse(x)), st
+    return __reverse(x), st
 end
 
 @inline function (r::ReverseSequence{Nothing})(
         x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N}
-    return reverse(x; dims=ndims(x) - 1), st
+    return __reverse(x; dims=ndims(x) - 1), st
 end
 
 @inline function (r::ReverseSequence)(x::AbstractVector{T}, ps, st::NamedTuple) where {T}
-    r.dim == 1 && return reverse(x), st
-    throw(DimensionMismatch(lazy"Cannot specify a dimension other than 1 for AbstractVector{T}"))
+    r.dim == 1 && return __reverse(x), st
+    throw(ArgumentError("Cannot specify a dimension other than 1 for AbstractVector{T}"))
 end
 
 @inline function (r::ReverseSequence)(
         x::AbstractArray{T, N}, ps, st::NamedTuple) where {T, N}
-    return reverse(x; dims=r.dim), st
+    return __reverse(x; dims=r.dim), st
 end
 
 """

src/layers/display.jl

@@ -53,6 +53,12 @@ function _printable_children(l::Union{PairwiseFusion, Parallel})
     return merge((; l.connection), children.layers)
 end
 _printable_children(l::SkipConnection) = (; l.connection, l.layers)
+function _printable_children(l::BidirectionalRNN)
+    merge_mode = l.model.connection isa Broadcast.BroadcastFunction ? l.model.connection.f :
+                 nothing
+    return (; merge_mode, forward_cell=l.model.layers.forward_rnn.cell,
+        backward_cell=l.model.layers.backward_rnn.rnn.cell)
+end
 
 _show_leaflike(x) = Functors.isleaf(x)  # mostly follow Functors, except for:
 _show_leaflike(x::AbstractExplicitLayer) = false

src/layers/recurrent.jl

@@ -1,5 +1,13 @@
 abstract type AbstractRecurrentCell{use_bias, train_state} <: AbstractExplicitLayer end
 
+const AbstractDebugRecurrentCell = Experimental.DebugLayer{
+    <:Any, <:Any, <:AbstractRecurrentCell}
+
+function ConstructionBase.constructorof(::Type{<:AbstractRecurrentCell{
+        use_bias, train_state}}) where {use_bias, train_state}
+    return AbstractRecurrentCell{use_bias, train_state}
+end
+
 # Fallback for vector inputs
 function (rnn::AbstractRecurrentCell)(x::AbstractVector, ps, st::NamedTuple)
     (y, carry), st_ = rnn(reshape(x, :, 1), ps, st)
@@ -82,16 +90,16 @@ automatically operate over a sequence of inputs.
 
     For some discussion on this topic, see https://github.com/LuxDL/Lux.jl/issues/472.
 """
-struct Recurrence{
-    R, C <: AbstractRecurrentCell, O <: AbstractTimeSeriesDataBatchOrdering} <:
-       AbstractExplicitContainerLayer{(:cell,)}
-    cell::C
-    ordering::O
+@concrete struct Recurrence{R} <: AbstractExplicitContainerLayer{(:cell,)}
+    cell <: Union{<:AbstractRecurrentCell, <:AbstractDebugRecurrentCell}
+    ordering <: AbstractTimeSeriesDataBatchOrdering
 end
 
+ConstructionBase.constructorof(::Type{<:Recurrence{R}}) where {R} = Recurrence{R}
+
 function Recurrence(cell; ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex(),
         return_sequence::Bool=false)
-    return Recurrence{return_sequence, typeof(cell), typeof(ordering)}(cell, ordering)
+    return Recurrence{return_sequence}(cell, ordering)
 end
 
 _eachslice(x::AbstractArray, ::TimeLastIndex) = _eachslice(x, Val(ndims(x)))
@@ -164,9 +172,8 @@ update the state with `Lux.update_state(st, :carry, nothing)`.
       + `cell`: Same as `cell`.
       + `carry`: The carry state of the `cell`.
 """
-struct StatefulRecurrentCell{C <: AbstractRecurrentCell} <:
-       AbstractExplicitContainerLayer{(:cell,)}
-    cell::C
+@concrete struct StatefulRecurrentCell <: AbstractExplicitContainerLayer{(:cell,)}
+    cell <: Union{<:AbstractRecurrentCell, <:AbstractDebugRecurrentCell}
 end
 
 function initialstates(rng::AbstractRNG, r::StatefulRecurrentCell)
@@ -641,3 +648,69 @@ function Base.show(io::IO, g::GRUCell{use_bias, TS}) where {use_bias, TS}
     TS && print(io, ", train_state=true")
     return print(io, ")")
 end
+
+"""
+    BidirectionalRNN(cell::AbstractRecurrentCell,
+        backward_cell::Union{AbstractRecurrentCell, Nothing}=nothing;
+        merge_mode::Union{Function, Nothing}=vcat,
+        ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex())
+
+Bidirectional RNN wrapper.
+
+## Arguments
+
+  - `cell`: A recurrent cell. See [`RNNCell`](@ref), [`LSTMCell`](@ref), [`GRUCell`](@ref),
+    for how the inputs/outputs of a recurrent cell must be structured.
+  - `backward_cell`: A optional backward recurrent cell. If `backward_cell` is `nothing`,
+    the rnn layer instance passed as the `cell` argument will be used to generate the
+    backward layer automatically. `in_dims` of `backward_cell` should be consistent with
+    `in_dims` of `cell`
+
+## Keyword Arguments
+
+  - `merge_mode`: Function by which outputs of the forward and backward RNNs will be combined.
+    default value is `vcat`. If `nothing`, the outputs will not be combined.
+  - `ordering`: The ordering of the batch and time dimensions in the input. Defaults to
+    `BatchLastIndex()`. Alternatively can be set to `TimeLastIndex()`.
+
+## Inputs
+
+  - If `x` is a
+
+      + Tuple or Vector: Each element is fed to the `cell` sequentially.
+
+      + Array (except a Vector): It is spliced along the penultimate dimension and each
+        slice is fed to the `cell` sequentially.
+
+## Returns
+
+  - Merged output of the `cell` and `backward_cell` for the entire sequence.
+  - Update state of the `cell` and `backward_cell`.
+
+## Parameters
+
+  - `NamedTuple` with `cell` and `backward_cell`.
+
+## States
+
+  - Same as `cell` and `backward_cell`.
+"""
+@concrete struct BidirectionalRNN <: AbstractExplicitContainerLayer{(:model,)}
+    model <: Parallel
+end
+
+(rnn::BidirectionalRNN)(x, ps, st::NamedTuple) = rnn.model(x, ps, st)
+
+function BidirectionalRNN(cell::AbstractRecurrentCell,
+        backward_cell::Union{AbstractRecurrentCell, Nothing}=nothing;
+        merge_mode::Union{Function, Nothing}=vcat,
+        ordering::AbstractTimeSeriesDataBatchOrdering=BatchLastIndex())
+    layer = Recurrence(cell; return_sequence=true, ordering)
+    backward_rnn_layer = backward_cell === nothing ? layer :
+                         Recurrence(backward_cell; return_sequence=true, ordering)
+    fuse_op = merge_mode === nothing ? nothing : Broadcast.BroadcastFunction(merge_mode)
+    return BidirectionalRNN(Parallel(fuse_op;
+        forward_rnn=layer,
+        backward_rnn=Chain(;
+            rev1=ReverseSequence(), rnn=backward_rnn_layer, rev2=ReverseSequence())))
+end

src/utils.jl

@@ -375,3 +375,5 @@ end
 @inline __set_refval!(x, y) = (x[] = y)
 
 @inline __eltype(x) = eltype(x)
+
+@inline __reverse(x; dims=:) = reverse(x; dims)

test/layers/basic_tests.jl

@@ -35,7 +35,7 @@
 
             @test layer(x, ps, st)[1] == aType(xr)
             @test layer(x2, ps, st)[1] == aType(x2rd1)
-            @test_throws DimensionMismatch layer2(x, ps2, st2)[1]
+            @test_throws ArgumentError layer2(x, ps2, st2)[1]
             @test layer3(x, ps3, st3)[1] == aType(xr)
             @test layer2(x2, ps2, st2)[1] == aType(x2rd2)
 

test/layers/recurrent_tests.jl

@@ -374,3 +374,77 @@ end
         @test_throws ErrorException Lux._eachslice(x, BatchLastIndex())
     end
 end
+
+@testitem "Bidirectional" timeout=3000 setup=[SharedTestSetup] tags=[:recurrent_layers] begin
+    rng = StableRNG(12345)
+
+    @testset "$mode" for (mode, aType, device, ongpu) in MODES
+        @testset "cell: $_cell" for _cell in (RNNCell, LSTMCell, GRUCell)
+            cell = _cell(3 => 5)
+            bi_rnn = BidirectionalRNN(cell)
+            bi_rnn_no_merge = BidirectionalRNN(cell; merge_mode=nothing)
+            display(bi_rnn)
+
+            # Batched Time Series
+            x = randn(rng, Float32, 3, 4, 2) |> aType
+            ps, st = Lux.setup(rng, bi_rnn) .|> device
+            y, st_ = bi_rnn(x, ps, st)
+            y_, st__ = bi_rnn_no_merge(x, ps, st)
+
+            @jet bi_rnn(x, ps, st)
+            @jet bi_rnn_no_merge(x, ps, st)
+
+            @test size(y) == (4,)
+            @test all(x -> size(x) == (10, 2), y)
+
+            @test length(y_) == 2
+            @test size(y_[1]) == size(y_[2])
+            @test size(y_[1]) == (4,)
+            @test all(x -> size(x) == (5, 2), y_[1])
+
+            __f = p -> sum(Base.Fix1(sum, abs2), first(bi_rnn(x, p, st)))
+            @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu
+
+            __f = p -> begin
+                (y1, y2), st_ = bi_rnn_no_merge(x, p, st)
+                return sum(Base.Fix1(sum, abs2), y1) + sum(Base.Fix1(sum, abs2), y2)
+            end
+            @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu
+
+            @testset "backward_cell: $_backward_cell" for _backward_cell in (
+                RNNCell, LSTMCell, GRUCell)
+                cell = _cell(3 => 5)
+                backward_cell = _backward_cell(3 => 5)
+                bi_rnn = BidirectionalRNN(cell, backward_cell)
+                bi_rnn_no_merge = BidirectionalRNN(cell, backward_cell; merge_mode=nothing)
+                display(bi_rnn)
+
+                # Batched Time Series
+                x = randn(rng, Float32, 3, 4, 2) |> aType
+                ps, st = Lux.setup(rng, bi_rnn) .|> device
+                y, st_ = bi_rnn(x, ps, st)
+                y_, st__ = bi_rnn_no_merge(x, ps, st)
+
+                @jet bi_rnn(x, ps, st)
+                @jet bi_rnn_no_merge(x, ps, st)
+
+                @test size(y) == (4,)
+                @test all(x -> size(x) == (10, 2), y)
+
+                @test length(y_) == 2
+                @test size(y_[1]) == size(y_[2])
+                @test size(y_[1]) == (4,)
+                @test all(x -> size(x) == (5, 2), y_[1])
+
+                __f = p -> sum(Base.Fix1(sum, abs2), first(bi_rnn(x, p, st)))
+                @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu
+
+                __f = p -> begin
+                    (y1, y2), st_ = bi_rnn_no_merge(x, p, st)
+                    return sum(Base.Fix1(sum, abs2), y1) + sum(Base.Fix1(sum, abs2), y2)
+                end
+                @eval @test_gradients $__f $ps atol=1e-2 rtol=1e-2 gpu_testing=$ongpu
+            end
+        end
+    end
+end