orbax
orbax copied to clipboard
google
Orbax provides common utility libraries for JAX users.
Current implementation of `slice_count` is not what the name indicates. Actually it returns the **replica** count along the `replica_axis_index`. This PR returns the actual slice count based on the global...
**Changes** Improved dependency granularity by changing the step library to import only the necessary class from Phoenix, instead of a broader module. - Modifies Orbax to depend on this new,...
![copybara-service[bot] avatar](https://avatars.githubusercontent.com/in/44061?v=4 "copybara-service[bot] avatar"){kind=avatar}
I'm trying to migrate from an older version of orbax (`v0.6.4`) to a newer one (`v0.11.10`), and I have a problem with restoring checkpoints. I have two similar checkpoints, and...
Divide values into chunks when communicating with JAX distributed client
internal monitoring and hashing change