orbax

orbax copied to clipboard

Introduce `ListKey` and `TupleKey` for `list` and `tuple` PyTree nodes respectively.

copybara-service[bot]

Add Metadata Checkpointer for PWoC.

1

Add Metadata Checkpointer for PWoC.

copybara-service[bot]

Support nested `asyncio.run` without `nest_asyncio` library.

copybara-service[bot]

Add monitoring to record how long directory creation takes (per sequential instance).

copybara-service[bot]

Add `blocking_metadata_write` option to allow disabling async metadata write behavior.

copybara-service[bot]

Restore from local with the mutated mesh, and transfer arrays to get back Pytree with original mesh before broadcasting.

1

Restore from local with the mutated mesh, and transfer arrays to get back Pytree with original mesh before broadcasting.

copybara-service[bot]

I have been trying to use Orbax for checkpointing Flax NNX models and getting checkpointing to work for models with Dropout layers which also hold JAX RNG keys is not...

prabhudavidsheryl

Add emergency checkpoint logging of arrays for debugging.

copybara-service[bot]

Bumps the pip group with 1 update in the /docs/requirements directory: [ipython](https://github.com/ipython/ipython). Updates `ipython` from 7.23.1 to 8.10.0 Commits 15ea1ed release 8.10.0 560ad10 DOC: Update what's new for 8.10 (#13939)...

dependabot[bot]

test restoring to global_mesh in parallel

copybara-service[bot]