flax

[nnx] RNN: add broadcast_rngs and state_axes APIs

#4407

Merged