Add a guide for writing custom recurrent layers

[seanmor5]

Closes #511.

Axon.Layers.dynamic_unroll/7 and static_unroll/7 are public, but their docs did not say what cell_fn is called with, what it must return, what carry and mask mean, or what the unroll returns. #511 asked for exactly that plus a guide that builds a recurrent layer from scratch instead of treating Axon.lstm as a black box, and shows how to deal with padded, variable-length sequences. The built-in RNNs are generated by a macro, which made it hard to see that the underlying pieces are ordinary functions anyone can use.

This is a documentation PR; no library behavior changes.

The guide

guides/model_creation/custom_recurrent_layers.livemd builds an Elman RNN in three steps and uses it to classify padded character sequences:

1. The cell contract. It states up front what an unroll calls the cell with (input, carry, mask, input_kernel, recurrent_kernel, bias) and what it expects back ({output, new_carry}), and that the cell, not the unroll, is responsible for honoring the mask (1 means padding).
2. Calling the unroll on raw tensors. Before any Axon graph is involved, the guide runs static_unroll, dynamic_unroll and an Enum.map_reduce/3 over the time axis and shows all three agree, so readers see the unroll is just a scan. It explains when to prefer static (inlined, short sequences) vs dynamic (while loop, graph size independent of length).
3. Wrapping it as a layer and masking. SimpleRNN.simple_rnn/3 mirrors Axon.lstm/4: it declares params with Axon.param/3, calls Axon.layer/3 with the sequence, an Axon.mask/3 node and the params, and splits the {outputs, {hidden}} result with Axon.elem/2. A synthetic, seedable "names" dataset (two made-up languages, words of 3..10 letters padded with 0) replaces the downloaded name files from the PyTorch tutorial the reporter was porting, so the notebook is self-contained. The guide then demonstrates that the final hidden state equals the state after the last real letter and does not change through the padding, that dropping the mask changes the result, and that unroll: :static gives the same numbers. It finishes with Axon.Loop training (~94% test accuracy) and a "going further" section covering stacking, Axon.bidirectional/4, richer carries, and what a cell can and cannot call.

The core of it:

defn cell(input, {hidden}, mask, input_kernel, recurrent_kernel, bias) do
  candidate =
    Nx.tanh(
      Axon.Layers.dense(input, input_kernel, bias) +
        Axon.Layers.dense(hidden, recurrent_kernel, 0)
    )

  mask = Nx.broadcast(Nx.as_type(mask, :u8), hidden)
  new_hidden = Nx.select(mask, hidden, candidate)
  {new_hidden, {new_hidden}}
end

tokens = Axon.input("tokens", shape: {nil, 10})
mask = Axon.mask(tokens, 0)

{sequence, hidden} =
  tokens
  |> Axon.embedding(27, 16)
  |> SimpleRNN.simple_rnn(32, mask: mask, name: "rnn")

Two things the guide calls out explicitly because they tripped up prototyping: Axon.param/3 shape functions receive one shape per Axon input of the layer (so a layer with a sequence and a mask input needs an arity-2 function, and the [{:axis, -1}, units] shape DSL cannot be used there), and Axon.mask/3 has to be computed from the integer tokens before Axon.embedding/4.

Doc changes

• Axon.Layers.dynamic_unroll/7 now documents the full cell_fn contract, every argument, the mask semantics, the return value, and the while-loop trade-off. static_unroll/7 refers to it and explains the inlining trade-off.
• Axon.param/3 finishes its truncated sentence about shape functions: they receive one shape per Axon input of the layer, in order, so a layer with two Axon inputs needs an arity-2 function.
• Axon.mask/3 says what it outputs (u8, 1 at eos_token), that it must come from the token input, and where to pass it.
• Removed the copy-pasted "More memory efficient than traditional LSTM" line from the lstm_cell doc.
• The guide is registered in mix.exs extras and guides/guides.md; it lands under "Guides: Model Creation" via the existing wildcard.

Limitations

• The guide uses Axon.elem/2, which is on main but not in the 0.8.1 release, so the notebook needs the next release (or a git dependency) to run. The Mix.install follows the other guides.
• Axon.bidirectional/4 works with the custom layer but only without a mask: the wrapped function runs inside an Axon.block/2, which cannot reference the mask node from the surrounding graph, and the backward pass would need a reversed mask anyway. The guide says so instead of claiming unqualified compatibility.
• The guide uses the default :glorot_uniform initializer for the recurrent kernel rather than Axon.Initializers.orthogonal/1, because the orthogonal initializer's QR host callback breaks init_fn with an Nx.template under EXLA. That is an existing issue and is not addressed here.

Tests

test/axon/recurrent_guide_test.exs runs the guide's code, following the convention of serialization_guide_test.exs, with the SimpleRNN module verbatim and a smaller max_len = 8 dataset:

• static_unroll, dynamic_unroll and an Enum.map_reduce scan agree on outputs and final carry, with the expected shapes.
• The custom layer initializes input_kernel, recurrent_kernel and bias with the expected shapes under the "rnn" name.
• The mask freezes the state through padding (state at the last real token equals the final state and the state at the last padding step), the unmasked layer with the same params gives a different final state, and static and dynamic unrolls agree inside a model.
• Training with Axon.Loop.trainer on padded sequences returns a ModelState with the RNN params, loss decreases across epochs and accuracy exceeds 80% (deterministic via seed: 42).

The file runs in ~2s under the default Nx.Defn.Evaluator and passes under USE_EXLA=1 as well.

🤖 Generated with Claude Code