keras-nightly 3.12.0.dev2025100503__py3-none-any.whl → 3.14.0.dev2026011604__py3-none-any.whl

keras/src/saving/file_editor.py

@@ -455,6 +455,9 @@ class KerasFileEditor:
     def _extract_weights_from_store(self, data, metadata=None, inner_path=""):
         metadata = metadata or {}
 
+        # ------------------------------------------------------
+        # Collect metadata for this HDF5 group
+        # ------------------------------------------------------
         object_metadata = {}
         for k, v in data.attrs.items():
             object_metadata[k] = v
@@ -462,26 +465,98 @@ class KerasFileEditor:
             metadata[inner_path] = object_metadata
 
         result = collections.OrderedDict()
+
+        # ------------------------------------------------------
+        # Iterate over all keys in this HDF5 group
+        # ------------------------------------------------------
         for key in data.keys():
-            inner_path = f"{inner_path}/{key}"
+            # IMPORTANT:
+            # Never mutate inner_path; use local variable.
+            current_inner_path = f"{inner_path}/{key}"
             value = data[key]
+
+            # ------------------------------------------------------
+            # CASE 1 — HDF5 GROUP → RECURSE
+            # ------------------------------------------------------
             if isinstance(value, h5py.Group):
+                # Skip empty groups
                 if len(value) == 0:
                     continue
+
+                # Skip empty "vars" groups
                 if "vars" in value.keys() and len(value["vars"]) == 0:
                     continue
 
-            if hasattr(value, "keys"):
+                # Recurse into "vars" subgroup when present
                 if "vars" in value.keys():
                     result[key], metadata = self._extract_weights_from_store(
-                        value["vars"], metadata=metadata, inner_path=inner_path
+                        value["vars"],
+                        metadata=metadata,
+                        inner_path=current_inner_path,
                     )
                 else:
+                    # Recurse normally
                     result[key], metadata = self._extract_weights_from_store(
-                        value, metadata=metadata, inner_path=inner_path
+                        value,
+                        metadata=metadata,
+                        inner_path=current_inner_path,
                     )
-            else:
-                result[key] = value[()]
+
+                continue  # finished processing this key
+
+            # ------------------------------------------------------
+            # CASE 2 — HDF5 DATASET → SAFE LOADING
+            # ------------------------------------------------------
+
+            # Skip any objects that are not proper datasets
+            if not hasattr(value, "shape") or not hasattr(value, "dtype"):
+                continue
+
+            shape = value.shape
+            dtype = value.dtype
+
+            # ------------------------------------------------------
+            # Validate SHAPE (avoid malformed / malicious metadata)
+            # ------------------------------------------------------
+
+            # No negative dimensions
+            if any(dim < 0 for dim in shape):
+                raise ValueError(
+                    "Malformed HDF5 dataset shape encountered in .keras file; "
+                    "negative dimension detected."
+                )
+
+            # Prevent absurdly high-rank tensors
+            if len(shape) > 64:
+                raise ValueError(
+                    "Malformed HDF5 dataset shape encountered in .keras file; "
+                    "tensor rank exceeds safety limit."
+                )
+
+            # Safe product computation (Python int is unbounded)
+            num_elems = int(np.prod(shape))
+
+            # ------------------------------------------------------
+            # Validate TOTAL memory size
+            # ------------------------------------------------------
+            MAX_BYTES = 1 << 32  # 4 GiB
+
+            size_bytes = num_elems * dtype.itemsize
+
+            if size_bytes > MAX_BYTES:
+                raise ValueError(
+                    f"HDF5 dataset too large to load safely "
+                    f"({size_bytes} bytes; limit is {MAX_BYTES})."
+                )
+
+            # ------------------------------------------------------
+            # SAFE — load dataset (guaranteed ≤ 4 GiB)
+            # ------------------------------------------------------
+            result[key] = value[()]
+
+        # ------------------------------------------------------
+        # Return final tree and metadata
+        # ------------------------------------------------------
         return result, metadata
 
     def _generate_filepath_info(self, rich_style=False):

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/saving/saving_lib.py

@@ -943,7 +943,7 @@ class DiskIOStore:
         if self.archive:
             self.tmp_dir = get_temp_dir()
             if self.mode == "r":
-                self.archive.extractall(path=self.tmp_dir)
+                file_utils.extract_open_archive(self.archive, self.tmp_dir)
             self.working_dir = file_utils.join(
                 self.tmp_dir, self.root_path
             ).replace("\\", "/")

keras/src/testing/__init__.py

@@ -3,3 +3,4 @@ from keras.src.testing.test_case import jax_uses_gpu
 from keras.src.testing.test_case import tensorflow_uses_gpu
 from keras.src.testing.test_case import torch_uses_gpu
 from keras.src.testing.test_case import uses_gpu
+from keras.src.testing.test_case import uses_tpu

keras/src/testing/test_case.py

@@ -40,7 +40,20 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
         self.addCleanup(lambda: shutil.rmtree(temp_dir))
         return temp_dir
 
-    def assertAllClose(self, x1, x2, atol=1e-6, rtol=1e-6, msg=None):
+    def assertAllClose(
+        self,
+        x1,
+        x2,
+        atol=1e-6,
+        rtol=1e-6,
+        tpu_atol=None,
+        tpu_rtol=None,
+        msg=None,
+    ):
+        if tpu_atol is not None and uses_tpu():
+            atol = tpu_atol
+        if tpu_rtol is not None and uses_tpu():
+            rtol = tpu_rtol
         if not isinstance(x1, np.ndarray):
             x1 = backend.convert_to_numpy(x1)
         if not isinstance(x2, np.ndarray):
@@ -57,7 +70,9 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
             f"The two values are close at all elements. \n{msg}.\nValues: {x1}"
         )
 
-    def assertAlmostEqual(self, x1, x2, decimal=3, msg=None):
+    def assertAlmostEqual(self, x1, x2, decimal=3, tpu_decimal=None, msg=None):
+        if tpu_decimal is not None and uses_tpu():
+            decimal = tpu_decimal
         msg = msg or ""
         if not isinstance(x1, np.ndarray):
             x1 = backend.convert_to_numpy(x1)
@@ -195,6 +210,8 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
         run_training_check=True,
         run_mixed_precision_check=True,
         assert_built_after_instantiation=False,
+        tpu_atol=None,
+        tpu_rtol=None,
     ):
         """Run basic checks on a layer.
 
@@ -376,7 +393,9 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
                     msg="Unexpected number of torch_params",
                 )
 
-        def run_output_asserts(layer, output, eager=False):
+        def run_output_asserts(
+            layer, output, eager=False, tpu_atol=None, tpu_rtol=None
+        ):
             if expected_output_shape is not None:
 
                 def verify_shape(expected_shape, x):
@@ -422,7 +441,11 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
                         tree.flatten(expected_output), tree.flatten(output)
                     ):
                         self.assertAllClose(
-                            ref_v, v, msg="Unexpected output value"
+                            ref_v,
+                            v,
+                            msg="Unexpected output value",
+                            tpu_atol=tpu_atol,
+                            tpu_rtol=tpu_rtol,
                         )
                 if expected_num_losses is not None:
                     self.assertLen(layer.losses, expected_num_losses)
@@ -551,7 +574,13 @@ class TestCase(parameterized.TestCase, unittest.TestCase):
                 output_data = layer(**input_data, **call_kwargs)
             else:
                 output_data = layer(input_data, **call_kwargs)
-            run_output_asserts(layer, output_data, eager=True)
+            run_output_asserts(
+                layer,
+                output_data,
+                eager=True,
+                tpu_atol=tpu_atol,
+                tpu_rtol=tpu_rtol,
+            )
 
             if run_training_check:
                 run_training_step(layer, input_data, output_data)
@@ -621,6 +650,17 @@ def uses_gpu():
     return False
 
 
+def uses_tpu():
+    # Condition used to skip tests when using the TPU
+    try:
+        devices = distribution.list_devices()
+        if any(d.startswith("tpu") for d in devices):
+            return True
+    except AttributeError:
+        return False
+    return False
+
+
 def uses_cpu():
     devices = distribution.list_devices()
     if any(d.startswith("cpu") for d in devices):

keras/src/utils/backend_utils.py

@@ -3,6 +3,7 @@ import importlib
 import inspect
 import os
 import sys
+import warnings
 
 from keras.src import backend as backend_module
 from keras.src.api_export import keras_export
@@ -124,9 +125,22 @@ def set_backend(backend):
 
     Example:
 
-    ```python
-    keras.config.set_backend("jax")
-    ```
+    >>> import os
+    >>> os.environ["KERAS_BACKEND"] = "tensorflow"
+    >>>
+    >>> import keras
+    >>> from keras import ops
+    >>> type(ops.ones(()))
+    <class 'tensorflow.python.framework.ops.EagerTensor'>
+    >>>
+    >>> keras.config.set_backend("jax")
+    UserWarning: Using `keras.config.set_backend` is dangerous...
+    >>> del keras, ops
+    >>>
+    >>> import keras
+    >>> from keras import ops
+    >>> type(ops.ones(()))
+    <class 'jaxlib.xla_extension.ArrayImpl'>
 
     ⚠️ WARNING ⚠️: Using this function is dangerous and should be done
     carefully. Changing the backend will **NOT** convert
@@ -138,7 +152,7 @@ def set_backend(backend):
 
     This includes any function or class instance that uses any Keras
     functionality. All such code needs to be re-executed after calling
-    `set_backend()`.
+    `set_backend()` and re-importing all imported `keras` modules.
     """
     os.environ["KERAS_BACKEND"] = backend
     # Clear module cache.
@@ -159,3 +173,16 @@ def set_backend(backend):
                 module_name = module_name[module_name.find("'") + 1 :]
                 module_name = module_name[: module_name.find("'")]
                 globals()[key] = importlib.import_module(module_name)
+
+    warnings.warn(
+        "Using `keras.config.set_backend` is dangerous and should be done "
+        "carefully. Already-instantiated objects will not be converted. Thus, "
+        "any layers / tensors / etc. already created will no longer be usable "
+        "without errors. It is strongly recommended not to keep around any "
+        "Keras-originated objects instances created before calling "
+        "`set_backend()`. This includes any function or class instance that "
+        "uses any Keras functionality. All such code needs to be re-executed "
+        "after calling `set_backend()` and re-importing all imported `keras` "
+        "modules.",
+        stacklevel=2,
+    )