orbax issues

End support for legacy Pax formats with directory structure:

End support for legacy Pax formats with directory structure: ``` checkpoints/ checkpoint_0 # msgpack file corresponding to step ``` Users can manually restructure to the following: ``` checkpoints/ checkpoint_0/ checkpoint...

copybara-service[bot]

Removes the TF dependency on the JaxModule and pushes TF dependencies down to the tensorflow_export. The serving_config processing now also occurs in the tensorflow_export class as it is strictly specific to TensorFlow. TF is still needed in the ExportManager since the public interface allows properties for getting the internal TF module and serving signatures.

1

Removes the TF dependency on the JaxModule and pushes TF dependencies down to the tensorflow_export. The serving_config processing now also occurs in the tensorflow_export class as it is strictly specific...

copybara-service[bot]

Fix slice selection logic in emergency checkpointing, so that a slice with a complete local checkpoint (if it exists) is always chosen as the secondary slice.

copybara-service[bot]

Hi, I am trying to save checkpoints using the following code:

2

Hi, I am trying to save checkpoints using the following code: ``` options = ocp.CheckpointManagerOptions( max_to_keep=self.max_checkpoints, create=True, best_fn=best_loss, best_mode="min", ) self.checkpoint_manager = ocp.CheckpointManager( os.path.join(self._out_dir, "checkpoints"), options=options, item_names=("state", "metadata"), item_handlers={ "state":...

raresdolga