From 89f9bc41c72b3f9e6fb912a85fcf0f4e7c263cd3 Mon Sep 17 00:00:00 2001
From: RecML authors <recommendations-ml@google.com>
Date: Thu, 6 Aug 2026 09:45:09 -0700
Subject: [PATCH] Internal change

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

diff --git a/recml/core/utils/keras_utils.py b/recml/core/utils/keras_utils.py
index 31eb566..5f09da2 100644
--- a/recml/core/utils/keras_utils.py
+++ b/recml/core/utils/keras_utils.py
@@ -74,6 +74,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 +82,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 +94,11 @@ 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,
+                ),
+            ),
         ),
     )
 
@@ -430,6 +437,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 +445,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 +453,11 @@ 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,
+                ),
+            ),
         ),
     )