search

google / orbax

Orbax provides common checkpointing and persistence utilities for JAX users
https://orbax.readthedocs.io/
Apache License 2.0
307 stars 36 forks source link

Now that JAX supports donate on device_put, use that instead of the ugly device_put -> block -> delete routine #1224

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Now that JAX supports donate on device_put, use that instead of the ugly device_put -> block -> delete routine