From ff5cb0fdee25576cc6fae7bcb72486190eaee5b9 Mon Sep 17 00:00:00 2001
From: RecML authors <recommendations-ml@google.com>
Date: Thu, 6 Aug 2026 14:58:55 -0700
Subject: [PATCH] Internal change

PiperOrigin-RevId: 960520363
---
 recml/core/utils/keras_utils.py | 20 ++++++++++++++++++++
 1 file changed, 20 insertions(+)

diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py
index 31eb566..d62a4ac 100644
--- a/recml/core/utils/keras_utils.py
+++ b/recml/core/utils/keras_utils.py
@@ -66,6 +66,10 @@ def _to_shape_dtype_struct(x: keras.Variable) -> jax.ShapeDtypeStruct:
   )
 
 
+def _is_cns_path(path: str) -> bool:
+  # copybara:insert return False
+
+
 class KerasOrbaxCheckpointManagerV2(ocp.CheckpointManager):
   """An Orbax checkpoint manager for Keras 3."""
 
@@ -74,6 +78,7 @@ def __init__(
       checkpoint_dir: str,
       max_to_keep: int = 5,
       save_interval_epochs: int = 1,
+      choose_store_cell: bool = True,
   ):
     """Initializes a KerasOrbaxCheckpointManager.
 
@@ -81,6 +86,7 @@ def __init__(
       checkpoint_dir: The directory to save checkpoints to.
       max_to_keep: The maximum number of checkpoints to keep.
       save_interval_epochs: The interval (in epochs) to save checkpoints.
+      choose_store_cell: Whether to dynamically select the CNS2 store cell.
     """
     if keras.backend.backend() != "jax":
       raise ValueError(
@@ -92,6 +98,12 @@ def __init__(
         options=ocp.CheckpointManagerOptions(
             save_interval_steps=save_interval_epochs,
             max_to_keep=max_to_keep,
+            file_options=ocp.options.FileOptions(
+                cns2_storage_options=ocp.options.Cns2StorageOptions(
+                    choose_store_cell=choose_store_cell
+                    and _is_cns_path(checkpoint_dir),
+                ),
+            ),
         ),
     )
 
@@ -430,6 +442,7 @@ def __init__(
       checkpoint_dir: str,
       max_to_keep: int = 5,
       save_interval_epochs: int = 1,
+      choose_store_cell: bool = True,
   ):
     """Initializes a KerasOrbaxCheckpointManager.
 
@@ -437,6 +450,7 @@ def __init__(
       checkpoint_dir: The directory to save checkpoints to.
       max_to_keep: The maximum number of checkpoints to keep.
       save_interval_epochs: The interval (in epochs) to save checkpoints.
+      choose_store_cell: Whether to dynamically select the CNS2 store cell.
     """
     super().__init__(
         directory=checkpoint_dir,
@@ -444,6 +458,12 @@ def __init__(
         options=ocp.CheckpointManagerOptions(
             save_interval_steps=save_interval_epochs,
             max_to_keep=max_to_keep,
+            file_options=ocp.options.FileOptions(
+                cns2_storage_options=ocp.options.Cns2StorageOptions(
+                    choose_store_cell=choose_store_cell
+                    and _is_cns_path(checkpoint_dir),
+                ),
+            ),
         ),
     )