RNN easy to get wrong? (with MWE)

[JoaoAparicio]

Problem: if I pass an LSTMCell an input with 2 dimensions, it trivially broadcasts over the outer dimension (and so both c and h have 2 dimensions). I expected RNN to do the same, and it can, but it was really easy to get it wrong and unnoticed.

Example. Input has 2 dimensions. The inner (feature) dimension is size 4 and the outer (broadcast) dimension is size 2. Hidden dimension is size 7. In a scenario where LSTM broadcasts trivially I would expect the hidden dimension to have shape (2,7), and that's the case.

inputs = jnp.array([jnp.arange(1,5), jnp.zeros((4,))])
# [[ 1.00,  2.00,  3.00,  4.00],
#  [ 0.00,  0.00,  0.00,  0.00]]
lstm = nnx.LSTMCell(in_features=4, hidden_features=7, rngs=rngs)
carry = lstm.initialize_carry((2,4), rngs=rngs)
c, h = carry
# c and h are all zeros, both shape (2,7)
carry, out = lstm(carry, inputs)
# out:
# [[-0.10,  0.22, -0.44, -0.01,  0.06,  0.04, -0.00],
#  [ 0.00,  0.00,  0.00,  0.00,  0.00,  0.00,  0.00]]

So far everything is as expected. Importantly, note that the second position of the broadcast dimension is all zeros, that's a sign broadcast is happening correctly.

Now lets see RNN. Take the same input, repeat it in the time dimension 3 times, so shape (time_steps, lstm_broadcast, features).

sequence_size = 3
inputs = jnp.array([jnp.array([jnp.arange(1,5), jnp.zeros((4,))]) for _ in range(sequence_size)])
# [[[ 1.00,  2.00,  3.00,  4.00],
#   [ 0.00,  0.00,  0.00,  0.00]],
#
#  [[ 1.00,  2.00,  3.00,  4.00],
#   [ 0.00,  0.00,  0.00,  0.00]],
#
#  [[ 1.00,  2.00,  3.00,  4.00],
#   [ 0.00,  0.00,  0.00,  0.00]]]

Run

lstm = nnx.LSTMCell(in_features=4, hidden_features=7, rngs=rngs)
rnn = nnx.RNN(lstm, unroll=sequence_size, rngs=rngs, return_carry=True)
(c, h), out = rnn(inputs)
out
# [[[ 0.58, -0.10, -0.15,  0.08,  0.00,  0.20,  0.29],
#   [ 0.22, -0.32, -0.20,  0.03,  0.09,  0.07,  0.14]],
#
#  [[ 0.58, -0.10, -0.15,  0.08,  0.00,  0.20,  0.29],
#   [ 0.22, -0.32, -0.20,  0.03,  0.09,  0.07,  0.14]],
#
#  [[ 0.58, -0.10, -0.15,  0.08,  0.00,  0.20,  0.29],
#   [ 0.22, -0.32, -0.20,  0.03,  0.09,  0.07,  0.14]]]

Now the presence of the broadcast dimension mixes the features. And furthermore the shape of c is wrong. It's (3,7) should be (2, 7). Something internally isn't aware that we would like the second dimension to be an LSTM broadcast.

Ah, but RNN has a major_time=True feature so lets try that:

lstm = nnx.LSTMCell(in_features=4, hidden_features=7, rngs=rngs)
rnn = nnx.RNN(lstm, unroll=sequence_size, rngs=rngs, return_carry=True, time_major=True)
(c, h), out = rnn(inputs)
out
#  [[[ 0.52, -0.42,  0.01, -0.74, -0.73, -0.36,  0.38],
#    [ 0.00,  0.00,  0.00,  0.00,  0.00,  0.00,  0.00]],
#
#   [[ 0.71, -0.57,  0.01, -0.89, -0.95, -0.44,  0.41],
#    [ 0.00,  0.00,  0.00,  0.00,  0.00,  0.00,  0.00]],
#
#   [[ 0.78, -0.64,  0.01, -0.94, -0.98, -0.46,  0.40],
#    [ 0.00,  0.00,  0.00,  0.00,  0.00,  0.00,  0.00]]]
c
#   [[ 1.43, -1.18,  1.84, -1.83, -2.89, -0.50,  0.52],
#    [ 0.00,  0.00,  0.00,  0.00,  0.00,  0.00,  0.00]]

Now the LSTM dimensions aren't mixed and the internal state has the right shape. Great.

Now what if we want to make this work with batch dimensions? (batch, time, lstm_broadcast, features). The problem now is that time is not longer the outer dimension. I did manage to make this work by stating that time is the outer dimension, but then wrapping in vmap for the batch dimension. However, it would be much nicer to have a solution like #4507 (comment) without having to manually vmap. Such solution does "just work" in the situation where LSTM takes only 1 dimension.

One possibility could be that instead of passing time_major: bool we could just pass time_axis: int. Currently this is inferred internally https://github.com/google/flax/blob/881685c75ebbfb44bf15c0161b11b4eedcfc455d/flax/nnx/nn/recurrent.py#L649C5-L649C14.

flax 0.10.2
jax 0.4.37
jaxlib 0.4.36
Python 3.11.11