flax
94f322a4 - Use jax.api.device_put_sharded() in place of private JAX apis.

Commit

5 years ago

Use jax.api.device_put_sharded() in place of private JAX apis. PiperOrigin-RevId: 332047658

References

Author

vanderplas

Committer

a-googler

a-googler

Parents

Loading