From b77bcd61a14e22dbe1d662e4449ca0d51473c6bf Mon Sep 17 00:00:00 2001 From: Daniel Ng Date: Wed, 16 Sep 2026 10:10:42 -0700 Subject: [PATCH] Internal Change. PiperOrigin-RevId: 982582515 --- .../_src/serialization/pathways_handler_registry.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/checkpoint/orbax/checkpoint/_src/serialization/pathways_handler_registry.py b/checkpoint/orbax/checkpoint/_src/serialization/pathways_handler_registry.py index bdfb31dc0..06e05916e 100644 --- a/checkpoint/orbax/checkpoint/_src/serialization/pathways_handler_registry.py +++ b/checkpoint/orbax/checkpoint/_src/serialization/pathways_handler_registry.py @@ -33,6 +33,8 @@ CheckpointingImpl = pathways_types.CheckpointingImpl +PATHWAYS_CONTROLLER_PROCESS_ID: str = 'pwcontroller' + @@ -79,14 +81,16 @@ def get_pathways_numpy_handler( ) -> type_handlers.NumpyHandler: """Returns the Pathways NumpyHandler.""" return type_handlers.NumpyHandler( - ocdbt_process_id='pwcontroller', + ocdbt_process_id=PATHWAYS_CONTROLLER_PROCESS_ID, deepcopy_host_arrays=deepcopy_host_arrays, ) def get_pathways_scalar_handler() -> type_handlers.ScalarHandler: """Returns the Pathways ScalarHandler.""" - return type_handlers.ScalarHandler(ocdbt_process_id='pwcontroller') + return type_handlers.ScalarHandler( + ocdbt_process_id=PATHWAYS_CONTROLLER_PROCESS_ID + ) def _register_numpy_and_scalar_handlers():