flax
97dde9ab - Implement jax_utils.replicate using jax.device_put_replicated

Commit

4 years ago

Implement jax_utils.replicate using jax.device_put_replicated PiperOrigin-RevId: 401169278

References

Author

jheek

Committer

a-googler

Parents