Handle PRNG keys in reshard by using jax.device_put.#120
Merged
copybara-service[bot] merged 1 commit intomainfrom Oct 3, 2025
Merged
Handle PRNG keys in `reshard` by using `jax.device_put`.#120copybara-service[bot] merged 1 commit intomainfrom
copybara-service[bot] merged 1 commit intomainfrom
Commits
Commits on Oct 3, 2025
- authored andcommitted