Bidirectional RNN (Flax)

Read a sequence in both directions

A forward recurrent network processes tokens from left to right. Its output at position t can depend on inputs 0 through t. A backward network processes the same sequence from right to left, so its aligned output at t can depend on inputs t through T - 1.

Run these two networks with independent parameters and concatenate their outputs at each original position. The result has shape (T, 2H): H forward features followed by H backward features.

Both halves include the current input. Together they can use context from the whole sequence, but a finite hidden vector need not preserve every detail. The backward half also needs later inputs, so this model cannot produce its complete representation immediately as tokens arrive in a left-to-right stream.

Reverse the input, then realign the output

This exercise uses two separate nn.RNN(nn.GRUCell(features=H)) instances. The wrapper adds a batch dimension of size one:

x                 (T, D)
x_b = x[None, ...] (1, T, D)
forward outputs   (1, T, H)
backward outputs  (1, T, H)
concatenated      (1, T, 2H)
remove batch      (T, 2H)

Call the forward RNN on x_b. Call the backward RNN on x_b[:, ::-1, :], then reverse that result on the time axis as well. The first reversal chooses processing order; the second puts the backward states alongside the inputs they belong to.

A simple cumulative sum makes the alignment visible without GRU weights:

import jax.numpy as jnp

x = jnp.array([1., 2., 4.])
forward = jnp.cumsum(x)                  # [1, 3, 7]
backward_processing = jnp.cumsum(x[::-1]) # [4, 6, 7]
backward = backward_processing[::-1]     # [7, 6, 4]
aligned = jnp.stack([forward, backward], axis=-1)

assert jnp.array_equal(aligned, jnp.array([[1., 7.], [3., 6.], [7., 4.]]))
assert jnp.array_equal(aligned.reshape(-1), jnp.array([1., 7., 3., 6., 7., 4.]))

This example illustrates order and concatenation only. A GRU computes learned gates rather than cumulative sums. Without the second reversal, row zero would pair the forward state for the first input with the backward state for the last input.

Use Linen's built-in GRU

Use nn.GRUCell defaults for the activations, initializers and zero initial carry. Linen's cell weights the previous state by its update gate and applies the reset gate after the recurrent candidate projection. This differs from the custom GRU cell exercise, which weights the candidate by its update gate and resets before that projection. Substituting the custom cell changes this problem's outputs.

Each direction shares its own cell parameters across time. The two directions do not share parameters with each other. Construct both unnamed RNN/cell pairs inside BiRNN in forward-then-backward order, preserving the built-in cells' parameter scopes and defaults.

nn.RNN supports inputs with zero or more batch dimensions; the leading size-one dimension is the layout used by this wrapper, not a general requirement of the API. With the default time_major=False, this wrapper's time axis is axis 1. The default return_carry=False returns all timestep outputs, not just a final state.

Every supplied timestep is valid. Do not supply sequence lengths, carry between calls, or padding masks. Each call starts both directions from zero carry. The model has no training step or output prediction head.