Fix RNN unrolling gradients and bidirectional LSTMs

[seanmor5]

RNN layers in Axon have been learning from the wrong gradients. When an RNN is unrolled the "dynamic" way, it becomes an Nx while loop. Nx figures out the gradient of a while loop by combining each step in the order the steps happened going forward. That shortcut only gives the right answer when the steps can be reordered without changing the result, and RNN steps cannot be reordered. So the gradients that came back out were simply wrong. Because :dynamic is the default, this hit every LSTM, GRU, and ConvLSTM unless someone went out of their way to pick :static. The fix is to remember the state going into each step on the way forward, then walk the sequence backward afterward and work out the gradient one step at a time, which is the order reverse-mode actually calls for.

To be sure this was a real bug and not just two versions disagreeing, I checked both against finite differences, which nudges a weight by a tiny amount and measures how much the loss moves. For one GRU weight the true answer was 0.09108. The static unroll said 0.09149, so it was right all along. The dynamic unroll said 0.40509, which is more than four times too big. After the fix, the dynamic unroll also says 0.09149. While rewriting that loop I also fixed a smaller problem where the output buffer was always created as f32, which made f64 models fail to build at all.

There were two more bugs in Axon.bidirectional, and both show up when you wrap an LSTM. A bidirectional layer runs the sequence forward and also backward, then has to flip the backward results around so they line up with the forward ones. The old code flipped every single thing that came back. That is correct for the per-step outputs, but an LSTM also hands back its final cell and hidden state, and those do not have a time dimension at all. Flipping them scrambled the hidden state along its feature dimension. Now only the leaves whose shape actually matches the input's time dimension get flipped, and the state is passed along untouched, which is what Keras does.

The other bidirectional bug is that both directions were built from a single Axon.block. A block shares its parameters everywhere it is used, so the forward and backward LSTMs were secretly sharing one set of weights instead of each learning their own. That quietly cut the layer's capacity in half. Each direction now gets its own block, named <name>_forward and <name>_backward. On the Keras IMDB model this brings the trainable parameter count from 2,658,945 up to 2,757,761, which is exactly the number Keras reports.

To check all of this end to end, this also adds an Axon port of the Keras bidirectional LSTM IMDB example, matching it layer for layer, along with a small Python script that converts the Keras IMDB data into a format Nx can read. It now trains at least as well as the Keras original:

            accuracy   loss   val_accuracy   val_loss
keras         0.9151  0.2263       0.8428      0.3650
this example  0.9201  0.3012       0.8644      0.3436

Getting there took one more piece of work that is not a bug fix. At first the example converged more slowly than Keras, and the giveaway was the training loss rather than the validation loss, which meant our model was fitting the training data less hard rather than generalizing worse. The cause was that Keras sets up its LSTM weights differently than Axon does. Keras keeps the input weights and the recurrent weights as one wide matrix each and cuts them into four pieces, one per gate, so the random draw sees the whole width at once. Axon keeps a separate tensor per gate and draws each one on its own, which makes the numbers a bit bigger than Keras' and means the four recurrent pieces are not orthogonal to each other. Keras also starts the forget gate's bias at one instead of zero, so that gate begins open and gradients can travel back through time right away. None of that can be expressed through Axon.lstm/3 today, so the example draws the wide matrix once and slices it, and sets the forget bias, rewriting the parameters before training starts. I checked the rewritten weights directly: the recurrent matrix is orthogonal to within 4.8e-7, the input weights match glorot over the full width, only the forget bias is one, and the values survive the training loop's own setup.

There are five new tests: gradient agreement between the two unrolls for GRU and for LSTM recurrent kernels, a finite-difference check anchoring the dynamic gradient to the true answer, an f64 check, and a test that the bidirectional state comes back un-reversed the way Keras does it. The full suite passes at 827 tests.

🤖 Generated with Claude Code

[polvalente]

LGTM, but please validate if I fixed things correctly!