Skip to content

Handle PRNG keys in reshard by using jax.device_put.#120

Merged
copybara-service[bot] merged 1 commit intomainfrom
test_814737436
Oct 3, 2025
Merged

Handle PRNG keys in `reshard` by using `jax.device_put`.#120
copybara-service[bot] merged 1 commit intomainfrom
test_814737436

Commits

Commits on Oct 3, 2025