orbax

orbax copied to clipboard

Custom TypeHandler and "No per-process OCDBT checkpoint subdirs" warning

4

Hello, I've recently created a custom type handler. Using it, and running on a single process I see the following warning ```python WARNING:absl:[process=0][thread=async_save_18] Skipping merge of OCDBT checkpoints: No per-process...

PhilipVinc

Add support to create dummy directories in orbax benchmark.

copybara-service[bot]

`step_prefix` cannot contain `_` -- Checkpoint manager does not recognized multiple `_`.

2

# Bug Description: When I created a checkpoint manager option like the following, ```py options = ocp.CheckpointManagerOptions(step_prefix="ppo_networks") with ocp.CheckpointManager( ".../model_checkpoints/7358284e-a603-453f-9024-f69a27a293c4", options=options, ) as mngr: mngr.restore(0) ``` with my directory looks...

scott-yj-yang

`ocp.tree.serialize_tree` filtering logic for sequences with empty leaves

1

Hi, I spotted ocp.tree.serialize_tree but it seems the serialization logic won't work if you have empty leaves within a sequence. This happens quite frequently with optax where you'll end up...

JesseFarebro

_validate_params fails on zero-sized arrays

3

Hi, @niketkumar @cpgaffney1, cc @dionhaefner The following attempts to serialize a zero-sized array, but it fails validation in `_validate_params`. I believe the problem is that `_validate_params` expects to find for...

hrbigelow

Code duplication around tensorstore_spec logic in orbax-checkpoint

2

Hi Orbax team, I was looking at Orbax code at the latest version==0.7.0 and found that pieces of code with quite heavy logic around tensorstore_spec creation seem to contain duplicates....

minotru

Internal.

copybara-service[bot]

Some internal changes.

copybara-service[bot]

Support single-slice checkpointing in `emergency.CheckpointManager`. Associated changes to adjust docstrings and `AbstractCheckpointManager` signatures.

copybara-service[bot]

Nothing here, but needed for diffbase

copybara-service[bot]