Training API support for initial/final recurrent state (TBPTT / chunked sequence training)

Hi, thank you for the excellent Mamba implementation. I am using Mamba for a real-time robot control policy. During inference, the model runs step by step and maintains a recurrent/convolutional state (hidden/cache state) continuously across the whole episode: obs_0 -> state_1 obs_1 + state_1 -> state_2 ... obs_t + state_t -> action_t, state_{t+1} However, the training API only accepts a full sequence and appears to initialize the recurrent state internally. It does not expose an initial_state input or a final_state output compatible with the inference cache/state API. This makes efficient chunked training difficult: Training on full episodes preserves the same continuous state behavior as inference, but episodes have highly variable lengths and cause significant multi-GPU load imbalance. Training on independent fixed-length chunks resets the state at every chunk boundary, creating a train–inference mismatch because inference state is never reset within an episode. Truncated BPTT would be a good solution, but it requires carrying the recurrent state between chunks while detaching it from the autograd graph. Would it be possible to support an API conceptually like: output, final_state = model( x, initial_state=state, return_final_state=True, ) where: initial_state uses the same representation as the incremental inference cache/state; final_state can be passed into the next training chunk; the caller can use final_state = detach(final_state) between chunks for TBPTT; the result is numerically consistent with processing the concatenated sequence in one forward pass, apart from floating-point differences. For example: state = None for chunk in episode_chunks: output, state = model( chunk, initial_state=state, return_final_state=True, ) loss = compute_loss(output) loss.backward() state = detach_state(state)

This would enable fixed-length chunk training, token/frame-balanced distributed batches, and better GPU utilization while retaining the continuous-state behavior used at inference. A clarification would also be very helpful if this is already supported through an existing lower-level API or cache object that I may have missed. Thanks again.