3.7. Recurrent Layers

RustyML ships 3 recurrent layers: SimpleRNN, LSTM, and GRU. They live under rustyml::neural_network::layers::recurrent, and the prelude re-exports them. All 3 share 1 contract. Each layer consumes a 3D sequence tensor, runs a recurrence over the time axis, and returns either the final hidden state or every hidden state.

If you know Keras, these layers carry the return_sequences and go_backwards flags with the same default of false. Each flag on its own means what it means in Keras. The 2 together do not. A layer with both set returns its states in processing order, so slot 0 holds the state that came from the LAST input timestep. Keras reverses that sequence back into input order.

Section 3.7.1 covers both flags, and 3.7.9 covers what to do about the pair. They carry no return_state, no bidirectional wrapper, and no in-cell dropout.

3.7.1. The input/output contract

Every recurrent layer expects a 3D input tensor with shape (batch_size, timesteps, features). By default it produces a 2D output (batch_size, units). The features axis is the width the build settled. This page writes it input_dim in the weight-shape formulas below. input_dim comes from the last axis of the shape that reaches the layer, never from an argument. units is the hidden width you set.

The recurrence consumes the timestep axis, so by default it does not appear in the output. Only the last hidden state, h_T, survives.

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::activation::Tanh;
use rustyml::neural_network::layers::recurrent::SimpleRNN;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    // units = 3 hidden neurons. The feature count comes from the build.
    let mut rnn = SimpleRNN::new(3, Tanh::new()).unwrap();
    rnn.build(&Shape::with_free_batch(&[2, 5, 4])).unwrap();

    // (batch = 2, timesteps = 5, features = 4)
    let x = Array::zeros((2, 5, 4)).into_dyn();

    // An inference context runs the recurrence and records nothing for a backward pass
    let mut ctx = Ctx::inference();
    let out = rnn.forward(&x, &mut ctx).unwrap();

    // The timestep axis is gone: only the last hidden state survives.
    println!("output shape: {:?}", out.shape()); // [2, 3] == (batch, units)
}

output shape: [2, 3]

The 3 layers share 1 constructor signature: new(units, activation) -> Result<Self, Error>. The feature count of a timestep is not an argument. Layer::build reads it from the last axis of the shape that reaches the layer, and it sizes the input kernel then.

The activation argument takes impl Into<Activation>. You can pass an Activation enum variant, for example Activation::Tanh, or a thin layer wrapper such as Tanh::new(), ReLU::new(), Sigmoid::new(), Linear::new(), or Softmax::new(). Every wrapper converts to the same enum. Softmax::new() carries the default axis -1. RustyML refuses a Softmax::new().with_axis(..) layer here, because an embedded softmax accepts that default alone. See 3.2. Dense Layers and Activations for the full activation catalog.

new returns Error::InvalidParameter when units is 0. RustyML has no return_state option, no bidirectional wrapper, and no in-cell dropout. Each layer uses 1 dense implementation.

2 builder methods change what the recurrence emits and which way it runs. Each one consumes and returns self, so you can chain them onto the constructor, and both default to false:

with_return_sequences(true) keeps the time axis. Slot k of that axis holds the state after step k. The last slot is exactly what the same layer returns with the flag off. The backward pass then expects a gradient of the same rank-3 shape. The cell state of an LSTM stays internal in both settings.

with_go_backwards(true) reverses the reading order alone. Processing step 0 consumes input timestep timesteps - 1. The layer does not turn the output back into input order. Slot 0 of a returned sequence holds the state that came from the last input timestep. This flag changes no shape. When you need the output back in input order, reverse the sequence yourself.

The 2 flags are independent, and the following program runs all 3 useful settings on 1 input:

use ndarray::Array;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::*;
use rustyml::neural_network::traits::UnaryLayer;
use rustyml::neural_network::Ctx;

fn main() {
    // 2 sequences, 5 timesteps, 4 features.
    let x = Array::from_shape_fn((2, 5, 4), |(n, t, f)| (n * 20 + t * 4 + f) as f32 * 0.01)
        .into_dyn();

    // 1 inference context serves all 3 passes, because it records nothing.
    let mut ctx = Ctx::inference();

    // The default returns the last hidden state only.
    let mut last_layer = SimpleRNN::new(3, Tanh::new()).unwrap().with_random_state(1);
    last_layer.build(&Shape::with_free_batch(&[2, 5, 4])).unwrap();
    let last = last_layer.forward(&x, &mut ctx).unwrap();
    assert_eq!(last.shape(), &[2, 3]);

    // return_sequences keeps the time axis, so the output is rank 3.
    let mut sequence_layer = SimpleRNN::new(3, Tanh::new())
        .unwrap()
        .with_random_state(1)
        .with_return_sequences(true);
    sequence_layer.build(&Shape::with_free_batch(&[2, 5, 4])).unwrap();
    let sequence = sequence_layer.forward(&x, &mut ctx).unwrap();
    assert_eq!(sequence.shape(), &[2, 5, 3]);

    // The last slot of the sequence is exactly the rank-2 output above.
    for n in 0..2 {
        for u in 0..3 {
            assert!((sequence[[n, 4, u]] - last[[n, u]]).abs() < 1e-6);
        }
    }

    // go_backwards reads the input from last step to first. The output stays in
    // processing order, so slot 0 holds the state that came from input step 4.
    let mut reversed_layer = SimpleRNN::new(3, Tanh::new())
        .unwrap()
        .with_random_state(1)
        .with_return_sequences(true)
        .with_go_backwards(true);
    reversed_layer.build(&Shape::with_free_batch(&[2, 5, 4])).unwrap();
    let reversed = reversed_layer.forward(&x, &mut ctx).unwrap();
    assert_eq!(reversed.shape(), &[2, 5, 3]);
    assert!((reversed[[0, 0, 0]] - sequence[[0, 0, 0]]).abs() > 1e-6);

    println!("last {:?}, sequence {:?}", last.shape(), sequence.shape());
}

A 2D or 4D input is a hard error, not a silent reshape. forward returns Error::InvalidInput for any input that is not 3D. forward and backward both take &self and a &mut Ctx, and a training pass parks its forward activations in that context. Calling backward before forward returns Error::NeuralNetwork(NnError::ForwardPassNotRun("SimpleRNN")), or "LSTM" or "GRU" for those layers. See 1.6. Error Handling for how these error variants work together.

Ctx::training() and Ctx::inference() pick the mode of a pass, and no layer of this page holds a training flag of its own. A layer that a caller drives by hand takes forward_mut instead, which builds the layer from the tensor and then completes the pass. SequentialBuilder::build already builds every layer of a model, so a model never needs that entry point.

RustyML differs from Keras in 1 way. The activation argument controls only the candidate or output nonlinearity, not the gates. In SimpleRNN, this activation applies to every hidden state. In LSTM and GRU, it applies to the candidate, and in LSTM, it also applies to the cell state before the output gate.

The gates always use sigmoid. RustyML has no separate recurrent_activation option like Keras has. The gate nonlinearity is fixed. Tanh is the default activation, and almost every published architecture uses it.

3.7.2. SimpleRNN and why gated cells exist

SimpleRNN is the textbook Elman recurrence. It starts from a zero hidden state. Each timestep mixes the current input with the previous hidden state, using 2 weight matrices and 1 bias:

h_0 = 0
h_t = activation( x_t @ W + h_{t-1} @ U + b )     for t = 1..T
output = h_T

Here, W is the input kernel (input_dim, units). U is the recurrent kernel (units, units). b is the bias (1, units). @ is a matmul that runs over the batch dimension.

RustyML initializes W with Xavier/Glorot uniform, and U with an orthogonal matrix (Gram-Schmidt). The orthogonal recurrent kernel is deliberate. It keeps the state transition norm-preserving at initialization. This is a cheap way to delay the vanishing-gradient problem described next.

That problem is vanishing gradients, and its counterpart, exploding gradients. Backpropagating from h_T to h_1 multiplies the upstream gradient by a fresh Jacobian at every step. That step is roughly grad_{t-1} = (activation'(h_t) * grad_t) @ U^T. Chaining T of those steps raises a matrix to the T-th power. If the matrix’s effective magnitude is below 1, the gradient decays exponentially, and the network cannot learn dependencies more than a few steps back. If the magnitude is above 1, the gradient explodes instead.

The orthogonal U keeps the U^T factor norm-preserving, and tanh' <= 1 keeps the product bounded. Even so, the product still tends toward zero over long sequences. This is why SimpleRNN works only for short sequences, up to a few dozen steps, and fails at long-range dependencies. LSTM and GRU exist to solve this problem.

3.7.3. LSTM: an additive memory highway

LSTM adds a second state, the cell state c_t. Its update is additive. LSTM uses 3 sigmoid gates to decide what to write, what to keep, and what to read. RustyML stores the 4 weight blocks fused side by side. The column order matches Keras, [input | forget | cell | output], written [i | f | g | o]:

i_t = sigmoid( x_t @ W_i + h_{t-1} @ U_i + b_i )     input gate  (how much candidate to write)
f_t = sigmoid( x_t @ W_f + h_{t-1} @ U_f + b_f )     forget gate (how much old cell to keep)
g_t = act(     x_t @ W_g + h_{t-1} @ U_g + b_g )     candidate   ("cell gate")
o_t = sigmoid( x_t @ W_o + h_{t-1} @ U_o + b_o )     output gate (how much cell to expose)
c_t = f_t * c_{t-1} + i_t * g_t                      cell state  (additive update)
h_t = o_t * act(c_t)                                 hidden state

Here * is elementwise multiplication. act is the configurable activation, Tanh by default, applied to both the candidate and the cell state. The key line is c_t = f_t * c_{t-1} + i_t * g_t. Its gradient with respect to c_{t-1} is just f_t, an elementwise multiply with no repeated matmul.

When the forget gate is open, meaning f is close to 1, the cell state carries the gradient backward with almost no loss. This is a near-identity highway, and the additive term feeds into it. Memory persists, and gradient flows, because the dominant path is addition, not repeated matrix multiplication.

RustyML initializes the forget-gate bias to 1.0, and every other bias to zero. This gives the memory highway a head start. Before training shapes the gates, f starts biased open, so the cell state, and its gradient, survive from the first epoch. The integration test lstm_forget_bias_is_one_not_zero checks this behavior. LSTM has 4 gates’ worth of parameters: param_count = 4 * (input_dim * units + units * units + units).

3.7.4. GRU: the same idea with merged gates

GRU keeps the additive-blend idea from LSTM. It folds the input and forget gates into a single update gate, and it drops the separate cell state. GRU needs only 3 weight blocks instead of 4. RustyML stores them fused, in the order [update | reset | candidate], written [z | r | h]. This order matches Keras:

z_t = sigmoid( x_t @ W_z + h_{t-1} @ U_z + b_z )          update gate
r_t = sigmoid( x_t @ W_r + h_{t-1} @ U_r + b_r )          reset gate
n_t = act( x_t @ W_h + (r_t * h_{t-1}) @ U_h + b_h )      candidate
h_t = z_t * h_{t-1} + (1 - z_t) * n_t                     hidden state

The update gate z_t runs a convex blend. When z is close to 1, the layer copies the previous hidden state through unchanged. This acts as a gradient highway, the same way a closed LSTM forget gate does. When z is close to 0, the layer replaces the previous state with the fresh candidate.

This is Keras’ convention. Some write-ups use the complement, so a z value ported from one of those needs flipping. The tests gru_update_gate_one_keeps_previous_hidden and gru_update_gate_zero_takes_the_candidate check these 2 extremes. The test gru_fused_kernel_first_block_is_the_update_gate checks the column order.

Note precisely where the reset gate acts. RustyML computes r_t * h_{t-1} before the candidate’s recurrent matmul, as (r_t * h_{t-1}) @ U_h. This matches the original Cho et al. formulation, which is Keras’ reset_after=False. It differs from the CuDNN reset_after=True variant, which applies the reset after the matmul and needs 2 biases per gate. RustyML uses a single bias per gate here.

GRU’s param_count = 3 * (input_dim * units + units * units + units). This is 3 quarters of an LSTM of the same width. In practice, GRU trains a little faster and matches LSTM on many tasks. LSTM sometimes performs better when a task needs a long, precisely controlled memory. Both layers initialize their input kernels with Xavier/Glorot, using the per-gate fan input_dim + units, not the fused width. Both layers initialize each gate’s recurrent block as an independent orthogonal matrix.

3.7.5. Weights, shapes, and setting them by hand

The trainable tensors and their shapes:

Fusing every gate into 1 matrix is not just cosmetic. It lets the input projection and the recurrent projection run as 1 large GEMM per timestep, instead of 1 GEMM per gate. This gives a real cache and SIMD benefit. Use weights() from the LayerBase trait to inspect the live arrays. All 3 layers give them the Keras names kernel, recurrent_kernel, and bias, and each entry carries a borrowed view:

use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::activation::Tanh;
use rustyml::neural_network::layers::recurrent::LSTM;
use rustyml::neural_network::traits::{LayerBase, UnaryLayer};

fn main() {
    // units = 8, and the build settles the 4 features per timestep.
    let mut lstm = LSTM::new(8, Tanh::new()).unwrap().with_random_state(42);
    lstm.build(&Shape::with_free_batch(&[2, 5, 4])).unwrap();

    // All 4 gates are fused side by side: width == 4 * units.
    for entry in lstm.weights() {
        println!("{:<16} {:?}", entry.name, entry.value.shape());
    }
}

kernel           [4, 32]
recurrent_kernel [8, 32]
bias             [1, 32]

Call with_random_state(seed) for reproducible initialization. On an unbuilt layer it records the seed, and the build spends it on the kernel and recurrent-kernel draws, in that fixed order. The forget-bias-1.0 rule still applies. Without a seed, RustyML seeds the weights from the global seed or from entropy. See 7.1. Reproducibility and Random Seeds.

To install weights by hand, for example when you port from another framework or unit-test an exact recurrence, every layer has set_weights(kernel, recurrent_kernel, bias). This method takes the fused matrices directly. LSTM and GRU also have set_gate_weights(...). It accepts 1 (kernel, recurrent_kernel, bias) triple per gate, and concatenates the triples into the fused layout for you.

The per-gate argument order is (input, forget, cell, output) for LSTM, 12 arrays in total. For GRU it is (reset, update, candidate), 9 arrays in total. Any shape mismatch returns Error::NeuralNetwork(NnError::WeightShape { .. }). When you save or load a whole model, these arrays serialize as-is, and each one keeps the name it has here. A recurrent layer at position 2 of a model therefore writes 2.kernel, 2.recurrent_kernel, and 2.bias. See 3.9. Saving and Loading Weights.

3.7.6. BPTT and the cost model

All 3 layers are the same layer over a different cell. 1 implementation holds the build, the batched input projection, the walk along the time axis, the cache, and the reductions below. The cell holds the arithmetic of 1 timestep, and that is the only part that tells the 3 apart.

Training uses backpropagation through time with a full unroll. There is no truncation window. A training pass parks everything the backward pass needs in the context.

Every layer parks the state that enters each timestep. For SimpleRNN and GRU that state is the hidden state, and for LSTM it is the hidden state and the cell state. Each cell then parks what only it reads. SimpleRNN parks nothing more. LSTM parks activation(c_t) and all 4 gate activations per timestep. GRU parks the reset and update gates, the candidate, and the r_t * h_{t-1} product.

An inference pass records none of it, and it therefore costs less than a training pass. Memory scales linearly with timesteps. Long sequences cost RAM, not just time. The layer itself holds none of these values, whatever the mode, so a built model serves several threads at once.

backward walks the timesteps in reverse. It expects an upstream gradient at the shape the forward pass gave back: (batch, units) by default, and (batch, timesteps, units) under return_sequences. At each step, backward computes the per-timestep pre-activation gradient. It threads grad_h, and grad_c for LSTM, back to the previous step. This part is inherently sequential over time.

backward then batches the weight-gradient reductions. It collapses the per-timestep dz values over (batch, timesteps), so the kernel and bias gradients each fall out of a single large GEMM. The recurrent-kernel gradient takes 1 GEMM per group of gate blocks that projected the same array. SimpleRNN and LSTM project the previous hidden state into every gate, so they hold 1 group and need 1 GEMM. GRU projects the previous hidden state into its update and reset blocks, and r_t * h_{t-1} into its candidate block. It holds 2 groups and needs 2 GEMMs.

The 3 results go into the gradient store of the context, under the names kernel, recurrent_kernel, and bias. backward returns the input gradient alone. The store sums the gradients that reach 1 address, so a layer that 2 positions of a graph model share receives the total. RustyML does not clip a gradient inside the layer. If you need gradient clipping, use an optimizer that offers it. See 3.4. Optimizers.

The performance shape is sequential over the time axis, and parallel over the batch and fused gate axis. The input projection x @ W does not depend on the recurrence. RustyML computes it once, as a single batched GEMM across all timesteps, up front. Only the h_{t-1} @ U term must run step by step. Each of those steps is itself a batch-parallel GEMM.

GRU saves a little more work here. It fuses the reset and update recurrent projections into 1 GEMM, because both gates read h_{t-1}. Only the candidate’s recurrent projection stays separate, because its input r_t * h_{t-1} depends on the freshly computed reset gate.

Every matmul goes to the gemmkit backend (see 6.2. Matrix Multiplication). gemmkit sizes its own parallelism from the amount of work. Wide layers and large batches get thread parallelism automatically. Small layers stay serial, to avoid overhead. A timestep’s fused gate projection always applies its bias in the same pass as the product. Activation fusion into that pass happens only for SimpleRNN with ReLU, not for LSTM, GRU, or any other activation.

The practical result: more sequences, meaning a larger batch, parallelize well. More timesteps do not, because that axis is a serial dependency chain. 7.3. Performance Tuning and Parallelism covers the thread-pool settings.

3.7.7. Stacking and building a model

A recurrent layer returns only the last hidden state by default, a 2D (batch, units) tensor. You cannot feed that output straight into another recurrent layer. A recurrent layer needs a 3D (batch, timesteps, features) input, and it rejects a 2D tensor with Error::InvalidInput.

with_return_sequences(true) on the lower layer removes the problem. That layer then emits (batch, timesteps, units), which is the rank the next layer reads. The build threads that units in as the feature width of the next layer. Set the flag on every layer of a stack except the last one. The following program stacks 2 LSTM layers on the same task the next example uses:

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    // 4 sequences, 3 timesteps, 1 feature. Target = sum of the 3 scalars.
    let x = Array::from_shape_vec(
        (4, 3, 1),
        vec![0.1, 0.2, 0.1, 0.3, 0.1, 0.2, 0.0, 0.2, 0.2, 0.2, 0.2, 0.1],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![0.4, 0.6, 0.4, 0.5])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        // (batch, 3, 1) -> (batch, 3, 8): the time axis survives.
        .add(
            LSTM::new(8, Tanh::new())
                .unwrap()
                .with_random_state(42)
                .with_return_sequences(true),
        )
        // (batch, 3, 8) -> (batch, 8): the second layer consumes the time axis.
        .add(LSTM::new(8, Tanh::new()).unwrap().with_random_state(43))
        .add(Dense::new(1, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 200).unwrap();
    let pred = model.predict(&x).unwrap();
    assert_eq!(pred.shape(), &[4, 1]);
    println!("prediction shape: {:?}", pred.shape());
}

RepeatVector covers the other pattern, in which the lower layer keeps its default and returns 1 state. It turns a (batch, units) state into a (batch, n, units) sequence, so a second recurrent layer can read it. Every one of the n steps holds the same vector, so the second layer sees a constant context, not the per-timestep states of the first. That is the encoder-decoder pattern, in which a fixed context seeds the decoder. It is the pattern to use when the 2 sequences have different lengths. See 3.5. Convolutional Layers, section 3.5.8, for the layer itself.

The standard pattern uses a recurrent layer as a sequence encoder, followed by a dense head that maps the final state to your targets. This composes cleanly. The recurrent layer turns (batch, timesteps, features) into (batch, units), and Dense consumes exactly that 2D shape.

The example below learns a real sequence task: predict the sum of a length-3 scalar sequence. It uses an LSTM encoder, a Dense readout, Adam, and mean-squared error. It converges in a few hundred full-batch epochs. Recall from 3.1. The Sequential Model that fit runs 1 full-batch gradient step per epoch:

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    // 4 sequences, 3 timesteps, 1 feature. Target = sum of the 3 scalars.
    let x = Array::from_shape_vec(
        (4, 3, 1),
        vec![
            0.1, 0.2, 0.1, // sum 0.4
            0.3, 0.1, 0.2, // sum 0.6
            0.0, 0.2, 0.2, // sum 0.4
            0.2, 0.2, 0.1, // sum 0.5
        ],
    )
    .unwrap()
    .into_dyn();
    let y = Array::from_shape_vec((4, 1), vec![0.4, 0.6, 0.4, 0.5])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        // Recurrent feature extractor: (batch, 3, 1) -> (batch, 16)
        .add(LSTM::new(16, Tanh::new()).unwrap().with_random_state(42))
        // Read the last hidden state out to a single scalar.
        .add(Dense::new(1, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        MeanSquaredError::new(),
    );

    model.fit(&x, &y, 300).unwrap();

    let pred = model.predict(&x).unwrap();
    println!("target      : {:?}", y.as_slice().unwrap());
    println!("prediction  : {:?}", pred.as_slice().unwrap());
}

After 300 epochs, the 4 predictions land within a few percent of [0.4, 0.6, 0.4, 0.5]. The LSTM has learned to accumulate the sequence, and the dense head reads the accumulator out. Swap LSTM for GRU or SimpleRNN, and the same code still compiles and trains. On this short sequence, all 3 layer types converge. This is because the vanishing-gradient advantage of the gated cells shows up only on long sequences.

model.summary() prints each recurrent layer’s output shape as (None, units), and as (None, timesteps, units) under return_sequences, where timesteps is the extent the build threaded in. Build the model against a free time axis, and that axis prints as None too. For anything larger than a toy dataset, use fit_with_batches instead of fit, so each epoch takes several mini-batch steps and reshuffles the data.

3.7.8. Embedding: from indices to vectors

Every layer on this page reads a float tensor. Text does not arrive that way. A tokenizer gives you word indices such as [7, 42, 3], and an index is a name, not a quantity you can add. Embedding is the bridge. It holds a trainable table with 1 row per vocabulary entry, and it replaces each index with that row.

Embedding::new(input_dim, output_dim) sets the table shape, and it keeps both arguments. input_dim is the vocabulary size, not an axis of any tensor, so the largest usable index is input_dim - 1. That is why the shape of the input settles nothing here, and why this constructor lost no argument. output_dim is the width of 1 vector. The layer appends that width as a new trailing axis, so the output rank is always 1 more than the input rank.

The middle row is the one that matters here. It is exactly the (batch_size, timesteps, features) contract from 3.7.1, so an Embedding feeds a recurrent layer with no reshape between them.

A Tensor holds f32, so the indices arrive as floating-point values. The layer truncates each value toward zero and then checks it. 2.0 and 2.9 both select row 2, and -0.5 selects row 0. That is the rule the Keras cast to a whole number uses. An index outside 0..input_dim, and any non-finite value, gives an Error::InvalidInput that names the offending value.

Keras leaves those 2 cases to its backend. An out-of-range index gives NaN, and a negative index wraps to the end of the table.

The table is the only parameter, so param_count() reports input_dim * output_dim. It starts from a uniform draw over [-0.05, 0.05], which is the Keras default for this layer, not the Xavier/Glorot rule that Dense uses. Call .with_random_state(seed) to make the draw reproducible. The layer gives the table the Keras name embeddings. Read it with weight("embeddings"), and write it with set_weights(table). Both use the (input_dim, output_dim) shape.

The backward pass adds the upstream gradient into the row that each index selected. A row that 2 positions selected gets the sum of both, and a row that no index selected keeps a gradient of exactly 0. The gradient is still stored densely over the whole table, so 1 optimizer step costs input_dim * output_dim whatever the batch holds. A momentum term or a weight decay also moves a row whose gradient is 0, because both rules act on the parameter itself. Keras behaves the same way on its dense update path.

The example below is the standard text pipeline in miniature. It classifies a 4-token sequence, and it stacks Embedding, LSTM, and Dense:

use ndarray::Array;
use rustyml::prelude::*;

fn main() {
    // 4 sentences of 4 token indices each, drawn from a 10-word vocabulary
    let x = Array::from_shape_vec(
        (4, 4),
        vec![
            1.0, 2.0, 3.0, 4.0, // opens with token 1
            5.0, 6.0, 7.0, 8.0, // opens with token 5
            1.0, 2.0, 7.0, 8.0, // opens with token 1
            5.0, 6.0, 3.0, 4.0, // opens with token 5
        ],
    )
    .unwrap()
    .into_dyn();
    // Label 1 when the sentence opens with token 1, and 0 when it opens with token 5
    let y = Array::from_shape_vec((4, 1), vec![1.0, 0.0, 1.0, 0.0])
        .unwrap()
        .into_dyn();

    let mut model = SequentialBuilder::new()
        // Indices to vectors: (batch, 4) -> (batch, 4, 8)
        .add(Embedding::new(10, 8).unwrap().with_random_state(42))
        // Sequence encoder: (batch, 4, 8) -> (batch, 16)
        .add(LSTM::new(16, Tanh::new()).unwrap().with_random_state(7))
        // Read the last hidden state out to 1 probability
        .add(Dense::new(1, Sigmoid::new()).unwrap().with_random_state(3))
        .build(&Shape::with_free_batch(x.shape()))
        .unwrap();
    model.compile(
        Adam::new(0.01, 0.9, 0.999, 1e-8, 0.0).unwrap(),
        BinaryCrossEntropy::new(),
    );

    model.summary();
    model.fit(&x, &y, 200).unwrap();

    let pred = model.predict(&x).unwrap();
    println!("target     : {:?}", y.as_slice().unwrap());
    println!("prediction : {:?}", pred.as_slice().unwrap());
}

After 200 epochs the 4 predictions sit within a thousandth of [1.0, 0.0, 1.0, 0.0]. model.summary() prints the embedding output shape as (None, 4, 8), and it needs no forward pass. The build handed the layer the (None, 4) shape of the index batch, and compute_output_shape appends the vector width to it.

1 Keras argument has no counterpart here. mask_zero marks index 0 as padding and builds a mask for the layers after it. RustyML propagates no mask, so the flag would change nothing. RustyML omits the flag rather than accepting and ignoring it. Pad your sequences to a fixed length, and let the model learn that the padding index carries no information.

3.7.9. A bidirectional model

There is no Bidirectional wrapper. A graph model builds one from parts, and the parts are 2 recurrent layers, a Reverse, and a merge layer:

let input = builder.input(Shape::known(&[batch, steps, features]));
let forward = builder.add(LSTM::new(8, Activation::Tanh)?.with_return_sequences(true), &[input]);
let backward = builder.add(
    LSTM::new(8, Activation::Tanh)?
        .with_return_sequences(true)
        .with_go_backwards(true),
    &[input],
);
let aligned = builder.add(Reverse::new(1), &[backward]);
let merged = builder.add(Concatenate::new(-1), &[forward, aligned]);

The Reverse is not optional. The backward branch emits its states in processing order, so its slot k holds the state of input timestep steps - 1 - k. The forward branch slot k holds timestep k instead. Merging them without the reversal pairs mismatched timesteps at every position, and nothing reports it.

Reverse::new(1) puts the branch back into input order. Slot 0 of the backward branch then holds the state that consumed input timestep 0, which is what the forward branch holds there. See 3.5. Convolutional Layers, section 3.5.8.

With return_sequences left at false, the reversal is unnecessary, and Reverse would refuse the rank-2 input. The branch returns 1 state per sample, and that state already consumed input timestep 0.

Give the 2 branches different seeds. They are separate layers with separate weights, and 2 layers built from 1 seed draw identical weights.

Do not share 1 layer between the 2 branches. A graph may call 1 layer from several nodes, and that is real weight sharing. go_backwards is a field of the layer, not an argument of the call. A shared layer therefore runs the SAME direction at both nodes, and the second branch carries no backward information at all. The model builds, predicts and trains, and nothing reports it.

A padded batch needs care. This crate carries no mask. Padding that sits at the end of a sequence sits at the front of the reversed branch, so the backward recurrence starts on the padding. Reverse the values of a ragged batch yourself, per sample, before the model reads them.