keras-nightly 3.14.0.dev2026011204__py3-none-any.whl → 3.14.0.dev2026011304__py3-none-any.whl

keras/src/backend/tensorflow/numpy.py

@@ -2151,7 +2151,7 @@ def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
 
 def ndim(x):
     x = convert_to_tensor(x)
-    return x.ndim
+    return x.shape.rank
 
 
 def nonzero(x):

keras/src/callbacks/orbax_checkpoint.py

@@ -8,7 +8,6 @@ from keras.src.api_export import keras_export
 from keras.src.callbacks.monitor_callback import (
     MonitorCallback,  # For metric monitoring logic
 )
-from keras.src.utils.io_utils import print_msg
 from keras.src.utils.module_utils import ocp
 
 # Context and AsyncOptions are accessed through the lazy-loaded ocp module
@@ -62,6 +61,11 @@ class OrbaxCheckpoint(MonitorCallback):
     This callback saves the model's weights and optimizer state asynchronously
     using Orbax, allowing training to continue without blocking for I/O.
 
+    **Multi-host Support**: When running in a multi-host distributed training
+    environment with JAX backend, this callback automatically coordinates
+    checkpointing across all hosts to ensure consistency and proper
+    synchronization. Multi-host checkpointing is only supported on JAX.
+
     Example:
 
     ```python
@@ -92,10 +96,6 @@ class OrbaxCheckpoint(MonitorCallback):
         verbose: Verbosity mode, 0 or 1.
         save_best_only: if `save_best_only=True`, it only saves when the model
             is considered the "best" based on the monitored quantity.
-        save_weights_only: if `save_weights_only=True`, only the model's
-            weights will be saved. Otherwise, the full model state
-            (weights, non-trainable variables, optimizer state, and
-            metrics state) will be saved. Defaults to False.
         mode: one of {'auto', 'min', 'max'}. Used with `save_best_only`.
         save_freq: `'epoch'` or integer. Frequency to save checkpoints.
         max_to_keep: Integer, maximum number of recent checkpoints to keep.
@@ -112,7 +112,6 @@ class OrbaxCheckpoint(MonitorCallback):
         monitor="val_loss",
         verbose=0,
         save_best_only=False,
-        save_weights_only=False,
         mode="auto",
         save_freq="epoch",
         initial_value_threshold=None,
@@ -129,7 +128,6 @@ class OrbaxCheckpoint(MonitorCallback):
         self.directory = directory
         self.verbose = verbose
         self.save_best_only = save_best_only
-        self.save_weights_only = save_weights_only
         self.save_freq = save_freq
         self.max_to_keep = max_to_keep
         self.save_on_background = save_on_background
@@ -138,6 +136,9 @@ class OrbaxCheckpoint(MonitorCallback):
         self._current_epoch = 0  # Keep track of epoch
         self._total_batches_seen = 0  # Global batch counter for step tracking
 
+        # Multi-host support
+        self._multihost_initialized = self._is_multihost_initialized()
+
         if self.save_freq != "epoch" and not isinstance(self.save_freq, int):
             raise ValueError(
                 f"Unrecognized save_freq: {self.save_freq}. "
@@ -151,14 +152,18 @@ class OrbaxCheckpoint(MonitorCallback):
                 ocp.training.preservation_policies.LatestN(max_to_keep)
             )
 
-        # Use AnyPreservationPolicy to combine them.
+        # Use AnyPreservationPolicy to combine them, or use directly
+        # if single policy
         preservation_policy = None
         if policies:
-            preservation_policy = (
-                ocp.training.preservation_policies.AnyPreservationPolicy(
-                    policies
+            if len(policies) == 1:
+                preservation_policy = policies[0]
+            else:
+                preservation_policy = (
+                    ocp.training.preservation_policies.AnyPreservationPolicy(
+                        policies
+                    )
                 )
-            )
 
         # Create the V1 Checkpointer with direct parameter passing
         # Orbax will handle directory creation on all processes as needed
@@ -167,6 +172,54 @@ class OrbaxCheckpoint(MonitorCallback):
             preservation_policy=preservation_policy,
         )
 
+    def _is_multihost_initialized(self):
+        """Check if multi-host environment is initialized."""
+        # Multi-host checkpointing is only supported on JAX backend
+        if backend.backend() != "jax":
+            return False
+
+        multihost = ocp.multihost
+        # Check if JAX distributed client is initialized
+        # (indicates multihost setup)
+        return multihost.is_jax_distributed_client_initialized()
+
+    def _sync_processes(self, key=None):
+        """Synchronize all processes across hosts."""
+        if not self._multihost_initialized:
+            return  # No-op for single host
+
+        multihost = ocp.multihost
+        sync_key = key or "orbax_checkpoint_sync"
+        multihost.sync_global_processes(sync_key)
+
+    def is_multihost_enabled(self):
+        """Return True if multi-host checkpointing is enabled and initialized.
+
+        This method can be used to check if the callback is operating in
+        a multi-host distributed training environment. Multi-host checkpointing
+        is only supported on JAX backend.
+
+        Returns:
+            bool: True if multi-host support is active, False otherwise.
+        """
+        return self._multihost_initialized
+
+    def is_primary_host(self):
+        """Return True if this process is the primary host in multi-host setup.
+
+        In multi-host environments, only the primary host typically handles
+        logging and coordination tasks. Multi-host checkpointing is only
+        supported on JAX backend.
+
+        Returns:
+            bool: True if this is the primary host, False otherwise.
+            Always returns True in single-host environments.
+        """
+        if not self._multihost_initialized:
+            return True  # Single host is always primary
+        multihost = ocp.multihost
+        return multihost.is_primary_host()
+
     def _should_save_on_batch(self, batch):
         """Check if we should save on this batch."""
         if self.save_freq == "epoch":
@@ -186,32 +239,14 @@ class OrbaxCheckpoint(MonitorCallback):
         return False
 
     def _save_checkpoint(self, step, logs=None):
-        """Save a checkpoint at the given step."""
+        """Save a checkpoint at the given step with multi-host coordination."""
 
         # --- Prepare Composite State (Backend-Agnostic) ---
         state_tree = _get_state_tree(self.model)
 
         # Save the nested state structures directly (preserving layer
         # names and structure)
-        if self.save_weights_only:
-            composite_state = {
-                "trainable_variables": state_tree["trainable_variables"],
-            }
-            if "non_trainable_variables" in state_tree:
-                composite_state["non_trainable_variables"] = state_tree[
-                    "non_trainable_variables"
-                ]
-        else:
-            composite_state = state_tree
-
-        # --- Save Logic (V1 API) ---
-        # All processes participate in distributed checkpointing
-        # Checkpointer is configured to save unconditionally when
-        # save_pytree is called
-        if self.verbose > 0:
-            print_msg(
-                f"OrbaxCheckpoint: Triggering async save for step {step}..."
-            )
+        composite_state = state_tree
 
         # Use a single with statement. If context_options is empty,
         # Context() uses defaults.
@@ -282,18 +317,16 @@ class OrbaxCheckpoint(MonitorCallback):
         except Exception:
             pass  # Ignore errors during cleanup
 
+        # Multi-host synchronization: ensure all hosts complete cleanup
+        self._sync_processes("checkpoint_cleanup")
+
     def wait_until_finished(self):
         """Wait for any in-progress checkpoint operations to complete.
         This method blocks until all asynchronous checkpoint save operations
-        have completed. It should be called before attempting to load
-        checkpoints if there might be pending save operations.
+        have completed across all hosts in a multi-host setup.
         """
-        # Wait for any async operations to complete
-        if hasattr(self.checkpointer, "wait"):
-            self.checkpointer.wait()
-        else:
-            # Fallback for older Orbax versions that don't have wait() method
-            while self.checkpointer.is_saving_in_progress():
-                import time
+        # Wait for any async operations to complete on this host
+        self.checkpointer.wait()
 
-                time.sleep(0.1)
+        # Multi-host synchronization: ensure all hosts complete
+        self._sync_processes("checkpoint_wait_complete")

keras/src/losses/losses.py

@@ -73,6 +73,14 @@ class MeanSquaredError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.0, 1.0])
+    >>> y_pred = keras.ops.array([0.9, 0.1, 0.8])
+    >>> loss = keras.losses.MeanSquaredError()
+    >>> loss(y_true, y_pred)
+    0.02
     """
 
     def __init__(
@@ -114,6 +122,14 @@ class MeanAbsoluteError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([1.0, 0.3, 1.0])
+    >>> y_pred = keras.ops.array([1.9, 0.3, 1.8])
+    >>> loss = keras.losses.MeanAbsoluteError()
+    >>> loss(y_true, y_pred)
+    0.5666667
     """
 
     def __init__(
@@ -155,6 +171,14 @@ class MeanAbsolutePercentageError(LossFunctionWrapper):
             `"float32"` unless set to different value
             (via `keras.backend.set_floatx()`). If a `keras.DTypePolicy` is
             provided, then the `compute_dtype` will be utilized.
+
+    Examples:
+
+    >>> y_true = keras.ops.array([100.0, 200.0, 300.0])
+    >>> y_pred = keras.ops.array([90.0, 210.0, 310.0])
+    >>> loss = keras.losses.MeanAbsolutePercentageError()
+    >>> loss(y_true, y_pred)
+    6.111111
     """
 
     def __init__(

keras/src/models/model.py

@@ -992,13 +992,18 @@ class Model(Trainer, base_trainer.Trainer, Layer):
                     self.non_trainable_variables, path_value_dict
                 )
             elif k == "optimizer_variables":
-                self._assign_variable_values(
-                    self.optimizer.variables, path_value_dict
-                )
+                if hasattr(self, "optimizer") and self.optimizer is not None:
+                    self._assign_variable_values(
+                        self.optimizer.variables, path_value_dict
+                    )
             elif k == "metrics_variables":
-                self._assign_variable_values(
-                    self.metrics_variables, path_value_dict
-                )
+                if (
+                    hasattr(self, "metrics_variables")
+                    and self.metrics_variables
+                ):
+                    self._assign_variable_values(
+                        self.metrics_variables, path_value_dict
+                    )
             else:
                 raise ValueError(f"Unknown variable name: {k}")
 

keras/src/ops/image.py

@@ -616,42 +616,98 @@ def extract_patches(
     padding="valid",
     data_format=None,
 ):
-    """Extracts patches from the image(s).
+    """Extracts patches from the image(s) or volume(s).
+
+    This function supports both 2D and 3D patch extraction based on the
+    `size` argument length, similar to how `keras.ops.conv` handles
+    different dimensions.
 
     Args:
-        images: Input image or batch of images. Must be 3D or 4D.
-        size: Patch size int or tuple (patch_height, patch_width)
-        strides: strides along height and width. If not specified, or
-            if `None`, it defaults to the same value as `size`.
-        dilation_rate: This is the input stride, specifying how far two
-            consecutive patch samples are in the input. For value other than 1,
-            strides must be 1. NOTE: `strides > 1` is not supported in
-            conjunction with `dilation_rate > 1`
+        images: Input image/volume or batch of images/volumes.
+            For 2D patches: 3D `(H, W, C)` or 4D `(N, H, W, C)`.
+            For 3D patches: 4D `(D, H, W, C)` or 5D `(N, D, H, W, C)`.
+        size: Patch size as int or tuple.
+            Length 2 tuple `(patch_height, patch_width)` or int for 2D patches.
+            Length 3 tuple `(patch_depth, patch_height, patch_width)` for
+            3D patches.
+        strides: Strides for patch extraction. If not specified, defaults
+            to `size` (non-overlapping patches).
+        dilation_rate: Dilation rate for patch extraction. Note that
+            `dilation_rate > 1` is not supported with `strides > 1`.
         padding: The type of padding algorithm to use: `"same"` or `"valid"`.
         data_format: A string specifying the data format of the input tensor.
             It can be either `"channels_last"` or `"channels_first"`.
-            `"channels_last"` corresponds to inputs with shape
-            `(batch, height, width, channels)`, while `"channels_first"`
-            corresponds to inputs with shape `(batch, channels, height, width)`.
-            If not specified, the value will default to
-            `keras.config.image_data_format`.
+            If not specified, defaults to `keras.config.image_data_format`.
 
     Returns:
-        Extracted patches 3D (if not batched) or 4D (if batched)
+        Extracted patches with shape depending on input and `size`:
+        - 2D patches: 3D (unbatched) or 4D (batched)
+        - 3D patches: 4D (unbatched) or 5D (batched)
 
     Examples:
 
+    >>> # 2D patches from batch of images
     >>> image = np.random.random(
     ...     (2, 20, 20, 3)
-    ... ).astype("float32") # batch of 2 RGB images
+    ... ).astype("float32")
     >>> patches = keras.ops.image.extract_patches(image, (5, 5))
     >>> patches.shape
     (2, 4, 4, 75)
-    >>> image = np.random.random((20, 20, 3)).astype("float32") # 1 RGB image
+
+    >>> # 2D patches from single image
+    >>> image = np.random.random((20, 20, 3)).astype("float32")
     >>> patches = keras.ops.image.extract_patches(image, (3, 3), (1, 1))
     >>> patches.shape
     (18, 18, 27)
+
+    >>> # 3D patches from batch of volumes
+    >>> volumes = np.random.random(
+    ...     (2, 10, 10, 10, 3)
+    ... ).astype("float32")
+    >>> patches = keras.ops.image.extract_patches(volumes, (3, 3, 3))
+    >>> patches.shape
+    (2, 3, 3, 3, 81)
+
+    >>> # 3D patches from single volume
+    >>> volume = np.random.random((10, 10, 10, 3)).astype("float32")
+    >>> patches = keras.ops.image.extract_patches(volume, (3, 3, 3))
+    >>> patches.shape
+    (3, 3, 3, 81)
     """
+    # Validate size argument
+    if not isinstance(size, int):
+        if not isinstance(size, (tuple, list)):
+            raise TypeError(
+                "Invalid `size` argument. Expected an int or a tuple. "
+                f"Received: size={size} of type {type(size).__name__}"
+            )
+        if len(size) not in (2, 3):
+            raise ValueError(
+                "Invalid `size` argument. Expected a tuple of length 2 or 3. "
+                f"Received: size={size} with length {len(size)}"
+            )
+
+    # Determine 2D vs 3D based on size argument
+    if not isinstance(size, int) and len(size) == 3:
+        # 3D patch extraction
+        if any_symbolic_tensors((images,)):
+            return ExtractPatches3D(
+                size=size,
+                strides=strides,
+                dilation_rate=dilation_rate,
+                padding=padding,
+                data_format=data_format,
+            ).symbolic_call(images)
+        return _extract_patches_3d(
+            images,
+            size,
+            strides,
+            dilation_rate,
+            padding,
+            data_format=data_format,
+        )
+
+    # 2D patch extraction (default)
     if any_symbolic_tensors((images,)):
         return ExtractPatches(
             size=size,

keras/src/saving/orbax_util.py

@@ -0,0 +1,26 @@
+"""Orbax checkpoint loading functionality."""
+
+import os
+
+from keras.src.utils.module_utils import ocp
+
+
+def is_orbax_checkpoint(filepath):
+    """Check if the given path is an Orbax checkpoint directory."""
+    if not os.path.exists(filepath):
+        return False
+
+    try:
+        return ocp.is_orbax_checkpoint(filepath)
+    except (ImportError, AttributeError):
+        # Fallback to check for orbax.checkpoint file if Orbax API not available
+        return os.path.isfile(os.path.join(filepath, "orbax.checkpoint"))
+
+
+def find_latest_orbax_checkpoint(checkpoint_dir):
+    """Find the latest checkpoint in an Orbax checkpoint directory."""
+    checkpointer = ocp.training.Checkpointer(directory=checkpoint_dir)
+    latest = checkpointer.latest
+    if latest is None:
+        raise ValueError(f"No valid checkpoints found in {checkpoint_dir}")
+    return os.path.join(checkpoint_dir, str(latest.step))

keras/src/saving/saving_api.py

@@ -6,13 +6,11 @@ from absl import logging
 from keras.src.api_export import keras_export
 from keras.src.legacy.saving import legacy_h5_format
 from keras.src.saving import saving_lib
+from keras.src.saving.orbax_util import find_latest_orbax_checkpoint
+from keras.src.saving.orbax_util import is_orbax_checkpoint
 from keras.src.utils import file_utils
 from keras.src.utils import io_utils
-
-try:
-    import h5py
-except ImportError:
-    h5py = None
+from keras.src.utils.module_utils import h5py
 
 
 @keras_export(["keras.saving.save_model", "keras.models.save_model"])
@@ -149,8 +147,6 @@ def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
         keras.layers.Softmax()])
     model.save("model.keras")
     loaded_model = keras.saving.load_model("model.keras")
-    x = np.random.random((10, 3))
-    assert np.allclose(model.predict(x), loaded_model.predict(x))
     ```
 
     Note that the model variables may have different name values
@@ -208,7 +204,7 @@ def load_model(filepath, custom_objects=None, compile=True, safe_mode=True):
     else:
         raise ValueError(
             f"File format not supported: filepath={filepath}. "
-            "Keras 3 only supports V3 `.keras` files and "
+            "Keras 3 only supports V3 `.keras` files, "
             "legacy H5 format files (`.h5` extension). "
             "Note that the legacy SavedModel format is not "
             "supported by `load_model()` in Keras 3. In "
@@ -288,15 +284,16 @@ def load_weights(model, filepath, skip_mismatch=False, **kwargs):
             objects_to_skip=objects_to_skip,
         )
     elif filepath_str.endswith(".h5") or filepath_str.endswith(".hdf5"):
-        if not h5py:
-            raise ImportError(
-                "Loading a H5 file requires `h5py` to be installed."
-            )
         if objects_to_skip is not None:
             raise ValueError(
                 "`objects_to_skip` only supports loading '.weights.h5' files."
                 f"Received: {filepath}"
             )
+        if not h5py.available:
+            raise ImportError(
+                "Loading HDF5 files requires the h5py package. "
+                "You can install it via `pip install h5py`"
+            )
         with h5py.File(filepath, "r") as f:
             if "layer_names" not in f.attrs and "model_weights" in f:
                 f = f["model_weights"]
@@ -308,9 +305,35 @@ def load_weights(model, filepath, skip_mismatch=False, **kwargs):
                 legacy_h5_format.load_weights_from_hdf5_group(
                     f, model, skip_mismatch
                 )
+    elif is_orbax_checkpoint(filepath):
+        # Load weights from Orbax checkpoint
+        from keras.src.utils.module_utils import ocp
+
+        filepath = str(filepath)
+
+        # Determine if this is a root directory or a step directory
+        items = os.listdir(filepath)
+        has_step_subdirs = any(
+            os.path.isdir(os.path.join(filepath, item)) and item.isdigit()
+            for item in items
+        )
+
+        if has_step_subdirs:
+            # It's a root directory, find the latest checkpoint
+            checkpoint_path = find_latest_orbax_checkpoint(filepath)
+        else:
+            # It's a step directory, use it directly
+            checkpoint_path = filepath
+
+        # Load checkpoint
+        loaded_state = ocp.load_pytree(checkpoint_path)
+
+        # Set the model state directly from the loaded state
+        model.set_state_tree(loaded_state)
     else:
         raise ValueError(
             f"File format not supported: filepath={filepath}. "
-            "Keras 3 only supports V3 `.keras` and `.weights.h5` "
-            "files, or legacy V1/V2 `.h5` files."
+            "Keras 3 only supports V3 `.keras` files, "
+            "`.weights.h5` files, legacy H5 format files "
+            "(`.h5` extension), or Orbax checkpoints."
         )

keras/src/utils/module_utils.py

@@ -44,15 +44,26 @@ class OrbaxLazyModule(LazyModule):
         try:
             parent_module = importlib.import_module("orbax.checkpoint")
             self.module = parent_module.v1
+            self.parent_module = parent_module
         except ImportError:
             raise ImportError(self.import_error_msg)
 
+    def __getattr__(self, name):
+        if name == "_api_export_path":
+            raise AttributeError
+        if self.module is None:
+            self.initialize()
+        if name == "multihost":
+            return self.parent_module.multihost
+        return getattr(self.module, name)
+
 
 tensorflow = LazyModule("tensorflow")
 gfile = LazyModule("tensorflow.io.gfile", pip_name="tensorflow")
 tensorflow_io = LazyModule("tensorflow_io")
 scipy = LazyModule("scipy")
 jax = LazyModule("jax")
+h5py = LazyModule("h5py")
 torch_xla = LazyModule(
     "torch_xla",
     import_error_msg=(

keras/src/version.py

@@ -1,7 +1,7 @@
 from keras.src.api_export import keras_export
 
 # Unique source of truth for the version number.
-__version__ = "3.14.0.dev2026011204"
+__version__ = "3.14.0.dev2026011304"
 
 
 @keras_export("keras.version")

{keras_nightly-3.14.0.dev2026011204.dist-info → keras_nightly-3.14.0.dev2026011304.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-nightly
-Version: 3.14.0.dev2026011204
+Version: 3.14.0.dev2026011304
 Summary: Multi-backend Keras
 Author-email: Keras team <keras-users@googlegroups.com>
 License: Apache License 2.0

{keras_nightly-3.14.0.dev2026011204.dist-info → keras_nightly-3.14.0.dev2026011304.dist-info}/RECORD

@@ -128,7 +128,7 @@ keras/regularizers/__init__.py,sha256=542Shphw7W8h4Dyf2rmqMKUECVZ8IVBvN9g1LWhz-b
 keras/saving/__init__.py,sha256=KvL2GZxjvgFgEhvEnkvqjIR9JSNHKz-NWZacXajsjLI,1298
 keras/src/__init__.py,sha256=Gi4S7EiCMkE03PbdGNpFdaUYySWDs_FcAJ8Taz9Y1BE,684
 keras/src/api_export.py,sha256=gXOkBOnmscV013WAc75lc4Up01-Kkg9EylIAT_QWctg,1173
-keras/src/version.py,sha256=0SYbNx7HTGlOGZmZu-eyZGV5m49SkT8Is9KNyZKk-bg,204
+keras/src/version.py,sha256=6Xtelgtl2s2j5h0H2_ulwbg1pYur71GFnblVZ9ACMxk,204
 keras/src/activations/__init__.py,sha256=0nL3IFDB9unlrMz8ninKOWo-uCHasTUpTo1tXZb2u44,4433
 keras/src/activations/activations.py,sha256=mogPggtp4CGldI3VOPNmesRxp6EbiR1_i4KLGaVwzL8,17614
 keras/src/applications/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -211,7 +211,7 @@ keras/src/backend/tensorflow/layer.py,sha256=69d40LwL4HhKRsCjj1VRpjfrQXXF8VV3vh0
 keras/src/backend/tensorflow/linalg.py,sha256=_lZVfdY1tFvrN7xwbt3INGoTR0yC5v-kI1Q0XppVibY,8773
 keras/src/backend/tensorflow/math.py,sha256=zTu_7Ff6B2Ro862z_xH0OCmIWbV74DjsO5UnfjYuOUQ,12370
 keras/src/backend/tensorflow/nn.py,sha256=6vtZHzUED6_blUPE1Tnc3GAxPpJ2ebxoaiMn80tTL9k,51328
-keras/src/backend/tensorflow/numpy.py,sha256=wIfpnlqab7HudnV-XPUG4MLBIKUQhrxj92yLgiw1ZKI,104617
+keras/src/backend/tensorflow/numpy.py,sha256=nIpMvr-g81I9KF74RD4AbU4e4t-0eFa9MND2Fh1u8Tk,104623
 keras/src/backend/tensorflow/optimizer.py,sha256=kFlyEOnGjEYdLpd8mpwhUeku78__xBfZbbrDWpJrq60,9307
 keras/src/backend/tensorflow/random.py,sha256=iO8V_soaDXZm9ewyAVbjudhsMj08C348c9Bz64nxXC4,6475
 keras/src/backend/tensorflow/rnn.py,sha256=99EJqbPdWddmG14zyjjhUZfU5zo9ObmslF_Mak7EmAs,34602
@@ -254,7 +254,7 @@ keras/src/callbacks/lambda_callback.py,sha256=q-nNr_k7MyYRP3HIetFsutcLkq78cUYxDD
 keras/src/callbacks/learning_rate_scheduler.py,sha256=II0SLxltUX3omRbGTYffd9KTWLRKtzW57SDRe70_t7E,2965
 keras/src/callbacks/model_checkpoint.py,sha256=Jt2mMKHKx0isrQnhiHADDOmwu72J594m93PBHy-zpV8,18570
 keras/src/callbacks/monitor_callback.py,sha256=-QBKqkKJ7Rg6L40Q80IScpvybmLoodLWcJoAgnTe_c4,4184
-keras/src/callbacks/orbax_checkpoint.py,sha256=cJCC620Ok8t0gyjt2U69wP8EXM9zdqYuf7L1oDxLOK0,11285
+keras/src/callbacks/orbax_checkpoint.py,sha256=hG_OClsm4lYQVTyCLWLJqgdYl6OXtFjz0J6z5eUtsmY,12529
 keras/src/callbacks/progbar_logger.py,sha256=BqddKoOyc8vxxtKriq5QD3n5JhVPUxkuWF2u1UlCriQ,3104
 keras/src/callbacks/reduce_lr_on_plateau.py,sha256=isJ9EzVo8jIu-_kWTFHpM_gaI5PbHTcUBM0keR9FRHA,4766
 keras/src/callbacks/remote_monitor.py,sha256=VDbNzCdddCDe_ZoeVvwV50oJkwOehhT_IDDYD8LzFOg,2727
@@ -476,7 +476,7 @@ keras/src/legacy/saving/saving_utils.py,sha256=8Sa2rmBGnTv86Tix20OgwF5vTLTpUYbfG
 keras/src/legacy/saving/serialization.py,sha256=hiwqO3Il861pkfm0Egaeph2XbhOlQQobmZjbZZgK32c,21368
 keras/src/losses/__init__.py,sha256=rt63Ye0f7YdAR0eV0EOj2J61DI6xNdp2ojonx6rB3wE,6595
 keras/src/losses/loss.py,sha256=8dCOv64yj9QC_GbcKT9M8YEC_Jr01wWuo-BBqFbfg0Q,8783
-keras/src/losses/losses.py,sha256=lVAuX3K4IzeRVvjvnejlieiuxtPRMvXtvmCrLZGsT9s,99534
+keras/src/losses/losses.py,sha256=MeFB4X3YLiTCw8sOEKpFUrSD4yv8E7hte91Gg-v04ok,100169
 keras/src/metrics/__init__.py,sha256=CydJsY38PR2lRN4irhO_wnlvgruTEAgSHp8eUYE0lwY,7410
 keras/src/metrics/accuracy_metrics.py,sha256=i_7ObnlyyE_UKDj8Nk5h5skakqpMlkMiphJ20eqcYho,18274
 keras/src/metrics/confusion_metrics.py,sha256=EKN1JGndT7pVesg_YAh8mGiM2wieAbGzXlw1ftuUGu4,62640
@@ -492,14 +492,14 @@ keras/src/metrics/regression_metrics.py,sha256=eLacV_8CKtzA26BJDJuncUDATuL1x8O6S
 keras/src/models/__init__.py,sha256=DPbBPSfIGgsufTfJH5U5xJOeN_Ef4FMadT7KKYg3Kjg,143
 keras/src/models/cloning.py,sha256=P0gMH3H9nyz6SMsdt4BQO05rXFa4qiqZk44rFpEnHsM,15945
 keras/src/models/functional.py,sha256=uD-qH9WwAUhaBrAEWAKnsVvKo0tvdHxa1M0dbBOE96Y,34086
-keras/src/models/model.py,sha256=BctaWtaNzMEK4ro14Lm1vDcK2ermkzBF2hMA_LlFGvk,41971
+keras/src/models/model.py,sha256=9kM6rbiAZOx3ycq2qM7QV6h2P1di57rA2HlljstSkh8,42215
 keras/src/models/sequential.py,sha256=CC9Q1BNB9m7TkgMHRyjOzhQvneng576wJpmdgHrACKY,14352
 keras/src/models/variable_mapping.py,sha256=FVtcgjBRqOxtvkzOE6kjG9SpcB9keDg2gS5LOTlXvG0,2181
 keras/src/ops/__init__.py,sha256=aORlvnrqY_eQl0EFLWdpHsXHnQ6JLSw1qhwJMr-VXJ0,644
 keras/src/ops/core.py,sha256=t06-MvptYb6ZVwmNj083JyUtzU4M6UTVXOT2vVHtKyU,42781
 keras/src/ops/einops.py,sha256=-pxW0_AzDQNsR7t2TJrzvYXBJpmLYA3fJoO0U_U96PY,6268
 keras/src/ops/function.py,sha256=QV9n1-xeTPDK_FJ3sjlHDWVH2sqDj96R6YQnJueMOlA,17821
-keras/src/ops/image.py,sha256=NAf68cwEmR2LJuPGF_N2mXFVUR1LjQLMFeq8rLKEtLw,66864
+keras/src/ops/image.py,sha256=-UBomIodLByXhxN9aSt0JUGn-N5yulXgRZ06XpXq8mM,68743
 keras/src/ops/linalg.py,sha256=3V8S_cgNxZZCIFcFj-FBHTdRqWNbimDtumMvfoc0f30,26736
 keras/src/ops/math.py,sha256=4qYMJ5qAPmeSyeF63YWoGbUkQt6f4_VX0enOChU4mXU,37233
 keras/src/ops/nn.py,sha256=04gjHB2BWusy4tWm59EO5Ns1paJC5umDNGwNCKzaJWQ,104658
@@ -543,7 +543,8 @@ keras/src/saving/__init__.py,sha256=vnrtfvnzW7Gwtxe5COhaMoEnVYB5iDe2YlqJ-DvqFIk,
 keras/src/saving/file_editor.py,sha256=tsUo9mQbMa8433tHTnOKWFhDeathYwDb0CeWcDTTTBQ,32089
 keras/src/saving/keras_saveable.py,sha256=aGIt1ajtsaamfUq18LM6ql8JEoQzi3HwzJEuwQ9bmKE,1285
 keras/src/saving/object_registration.py,sha256=OOO-7-SNfPoFkFsR_c5jzE6aSIDIlHlnMcm9IlI_Gbs,7357
-keras/src/saving/saving_api.py,sha256=hYMr0g_4zboDHUA4Dop7PVSPsGB0FBN7d29W4RhNPNI,11655
+keras/src/saving/orbax_util.py,sha256=0o05YKFjiePHgeW_d5fvuUAGzymqbJTeuquUR-7uVGE,906
+keras/src/saving/saving_api.py,sha256=PMkxXhtNNKX8GlwIsCP8-Plt19M012wNEk7i8BhxWzo,12670
 keras/src/saving/saving_lib.py,sha256=-uSXsojqzSl19FtW5FogCclvnu_nnVU3S-Si293DNq0,58723
 keras/src/saving/serialization_lib.py,sha256=yzCTm8hin__MGA2N5M5F-8Zbts5ZJVmINbrH4wEtIwI,30334
 keras/src/testing/__init__.py,sha256=7vVsV7Rn3rG99DdURgnH8ncpxagRwIE0uhH-R4qDyok,315
@@ -584,7 +585,7 @@ keras/src/utils/io_utils.py,sha256=Riv9TCCnz6xQLUvR1QC-UOCoGZ_KiNTwQVvLY6dKcX8,4
 keras/src/utils/jax_layer.py,sha256=ytws8NcxWzJ4kViBy3bc-Pk3st3_3L8RqXxgq9sYp1k,32912
 keras/src/utils/jax_utils.py,sha256=vY3P4S9mfWEjdirLd81ocKqeCm-UVfgQ1yTi6UHdBiM,322
 keras/src/utils/model_visualization.py,sha256=0ENeiq8q-qbyGjfcRixyyInb3aTxfcKCooKhZ1hSuI0,17794
-keras/src/utils/module_utils.py,sha256=r878q2qMRiudcfWw_OBFodM9g_qMtRkCX7xsWHB4uns,2531
+keras/src/utils/module_utils.py,sha256=FTZPMRLurURchLPX1tu-h3b-UoPW28faNOlDzpYDW6A,2894
 keras/src/utils/naming.py,sha256=bPowKBlgiVP_6XtVlNVHxrxheKuJy2c0e-oEM8ocZQY,1776
 keras/src/utils/numerical_utils.py,sha256=Uqe5nu1HXmiZuh5-MznomtDSVSO9FgFaltdDtGnN61o,7205
 keras/src/utils/progbar.py,sha256=Yg2Vp1xzqU7HnfDEGSeZsmOKAKYKA4oEHv7yAMaucYw,10358
@@ -614,7 +615,7 @@ keras/utils/bounding_boxes/__init__.py,sha256=jtvQll4u8ZY0Z96HwNhP1nxWEG9FM3gI-6
 keras/utils/legacy/__init__.py,sha256=oSYZz6uS8UxSElRaaJYWJEoweJ4GAasZjnn7fNaOlog,342
 keras/visualization/__init__.py,sha256=UKWmiy6sps4SWlmQi9WX8_Z53cPpLlphz2zIeHdwJpQ,722
 keras/wrappers/__init__.py,sha256=QkS-O5K8qGS7C3sytF8MpmO6PasATpNVGF8qtb7Ojsw,407
-keras_nightly-3.14.0.dev2026011204.dist-info/METADATA,sha256=R7TeLDjuyJduEDB21M8CmcsgMPQeVhCcy1zJxmqiAlk,6339
-keras_nightly-3.14.0.dev2026011204.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_nightly-3.14.0.dev2026011204.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
-keras_nightly-3.14.0.dev2026011204.dist-info/RECORD,,
+keras_nightly-3.14.0.dev2026011304.dist-info/METADATA,sha256=VvTqhhwTIlxRsKL9og2HY3vWDoIfXS2ITQ2jU8Mt2X4,6339
+keras_nightly-3.14.0.dev2026011304.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_nightly-3.14.0.dev2026011304.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
+keras_nightly-3.14.0.dev2026011304.dist-info/RECORD,,