orbax

orbax copied to clipboard

Add warning when GCS checkpoint is missing `commit_success.txt` (in case the checkpoint comes from not GCS originally).

rdyro

When importing the `orbax` checkpointer, our scripts throw warnings as below: ``` WARNING:absl:Type handler registry overriding type "" collision on scalar WARNING:absl:Type handler registry overriding type "" collision on scalar...

varunagrawal

Modernize typing in composite checkpoint handler.

1

Modernize typing in composite checkpoint handler.

copybara-service[bot]

Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.

1

Sort devices explicitly by process index, then id (as opposed to IDs alone). IDs may be randomly generated, and are not guaranteed to be ordered based on their process index.

copybara-service[bot]

Fix possible race condition in orbax root directory creation.

1

Fix possible race condition in orbax root directory creation.

copybara-service[bot]

Migrate jax dependency to >=4.30 for jax.export API.

copybara-service[bot]

Allow user-provided symbolic constraints

1

Allow user-provided symbolic constraints

copybara-service[bot]

Introduce replica_axis_index in init of emergency checkpoint manager

1

Introduce replica_axis_index in init of emergency checkpoint manager

copybara-service[bot]

test copybara, don't submit

copybara-service[bot]

Orbax: Refactor tensorstore spec construction and support local temporary index path

1

Orbax: Refactor tensorstore spec construction and support local temporary index path This separates the KvStore spec construction from the array spec construction, and integrates the options only used for writing...

copybara-service[bot]