orbax

orbax copied to clipboard

obm: Match TF output signature argument names. TF SavedModel picks output names based either on their position or the key in the enclosing dictionary. Orbax Model used to use only...

copybara-service[bot]

saving with default/global mesh is broken

2

On main and jax 0.7/0.8, if a global mesh is set with `jax.sharding.set_mesh` it is impossible to save a checkpoint because of the error shown below. Am I doing something...

PhilipVinc

obm: Use original argument names for TF input signatures. Previously, the original input names were substituted as "input_0", "input_1", etc. This change updates the logic to extract input argument names...

copybara-service[bot]

Disable `multiprocess-checkpoint-benchmarks` job due to limited chips

1

Disable `multiprocess-checkpoint-benchmarks` job due to limited chips

copybara-service[bot]

Run pytree_checkpoint_handler_test in github actions

1

Run pytree_checkpoint_handler_test in github actions

copybara-service[bot]

CheckpointManager randomly crashes the kernel of the JupyterLab

11

The following code in multiple cells in the JupyterLab crashes randomly when executed: ```python import jax import flax from flax import nnx import orbax.checkpoint as ocp print(jax.__version__) print(flax.__version__) print(ocp.__version__) path...

vfdev-5

Eliminate usages of `ParamInfo.path` in favor of `parent_dir / name`.

copybara-service[bot]

Add utilities for working with local steps on Pathways.

copybara-service[bot]

Add flag option to raise an error if Pathways client is initialized before _gemini_lost_callback is set.

copybara-service[bot]

Add multiprocess checkpoint benchmarks job.

1

Add multiprocess checkpoint benchmarks job.

copybara-service[bot]