flax
Exposing the experimental _split_transpose JAX scan parameter in Flax.
#3795
Merged
Merged