How to unroll a model end to end?

Hello

What's the correct way to unroll a model that contains an LSTM?

e.g. Suppose my model has 3 blocks from top to bottom:

```
spatial_block (not recurrent)
lstm
mlp (not recurrent)
```

I know how to unroll the LSTM N times, there's a module for that:

```
lstm_unrolled = nnx.RRN(lstm, unroll=N, ...)
```

But then, how do I do the same for the rest of the non-recurrent parts of my model?

It ocurred to me that `nnx.RNN` is perhaps generic enough to work with the full model end to end? The fact that it's annotated as taking a `nnx.RNNCellBase` suggests this is probably not the case.

Instead I tried using `nnx.RNN` to unroll the LSTM but then do the rest manually.
For example I managed to make this work (schematically):

```
spatialblock_w_timesteps = nnx.vmap(SpatialBlock.__call__, in_axes=(None, 0), out_axes=0)
lstm_unrolled = nnx.RNN(model.lstmcell, unroll=rnn_unroll, rngs=rngs)
mlp_w_timesteps = nnx.vmap(MLP.__call__, in_axes=(None, 0), out_axes=0)

x = spatialblock_w_timesteps(model.spatial_block, x)
x = lstm_unrolled(x)
x = mlp_w_timesteps(model.mlp, x)
```

This works. But because it has to be done by hand on a model by model bases, it scales badly for more complex models, and it's error prone.
Is there a way of unrolling the full model as easily as the LSTM component?