[FLAX] Pass pytrees to lax.psum directly, rather than calling tree_map(lax.psum, ...).
[FLAX] Pass pytrees to lax.psum directly, rather than calling tree_map(lax.psum, ...).

lax.psum supports pytree arguments, and by passing a pytree to lax.psum directly, JAX will generate a single tuple all-reduce input to XLA. This is important in multihost GPU, where compilation isn't completely deterministic as it is on TPU, so the client must ensure that the order of the all-reductions in the model is fully determined by the Python code.

PiperOrigin-RevId: 298888382

google/flax,google/flax