flax

934bfe8c - Flax RNG guide for jax.jit: clarify rng outputs are shared but not inputs.

Flax RNG guide for jax.jit: clarify rng outputs are shared but not inputs. Of course that's obvious, but only after one has fully embraced the SPMD mindset and abandoned the olden ways, so it doesn't hurt to say it. PiperOrigin-RevId: 808994300