flax

feat(nnx): add out_sharding to recurrent cells

#5255

Open