keras-nightly 3.12.0.dev2025083103__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/trainers/compile_utils.py

@@ -148,6 +148,7 @@ class CompileMetrics(metrics_module.Metric):
         self.built = False
         self.name = "compile_metrics"
         self.output_names = output_names
+        self._resolved_output_names = None
 
     @property
     def metrics(self):
@@ -175,10 +176,16 @@ class CompileMetrics(metrics_module.Metric):
 
     def build(self, y_true, y_pred):
         num_outputs = 1  # default
-        if self.output_names:
+        # Resolve output names. If y_pred is a dict, prefer its keys.
+        if isinstance(y_pred, dict):
+            keys = sorted(list(y_pred.keys()))
+            if self.output_names and set(self.output_names) == set(keys):
+                # If there is a perfect match, use the user-provided order.
+                output_names = self.output_names
+            else:
+                output_names = keys
+        elif self.output_names:
             output_names = self.output_names
-        elif isinstance(y_pred, dict):
-            output_names = sorted(list(y_pred.keys()))
         elif isinstance(y_pred, (list, tuple)):
             num_outputs = len(y_pred)
             if all(hasattr(x, "_keras_history") for x in y_pred):
@@ -187,6 +194,7 @@ class CompileMetrics(metrics_module.Metric):
                 output_names = None
         else:
             output_names = None
+        self._resolved_output_names = output_names
         if output_names:
             num_outputs = len(output_names)
 
@@ -316,9 +324,10 @@ class CompileMetrics(metrics_module.Metric):
         return flat_metrics
 
     def _flatten_y(self, y):
-        if isinstance(y, dict) and self.output_names:
+        names = self._resolved_output_names
+        if isinstance(y, dict) and names:
             result = []
-            for name in self.output_names:
+            for name in names:
                 if name in y:
                     result.append(y[name])
             return result
@@ -690,17 +699,34 @@ class CompileLoss(losses_module.Loss):
             return self.call(y_true, y_pred, sample_weight)
 
     def call(self, y_true, y_pred, sample_weight=None):
+        def resolve_path(path, object):
+            for _path in path:
+                object = object[_path]
+            return object
+
         if not tree.is_nested(y_true) and not tree.is_nested(y_pred):
             # Fast path: single output case / no loss-tracking metric.
             if not self.built:
                 self.build(y_true, y_pred)
-            _, loss_fn, loss_weight, _ = self._flat_losses[0]
-            loss_value = ops.cast(
-                loss_fn(y_true, y_pred, sample_weight), dtype=self.dtype
-            )
-            if loss_weight is not None:
-                loss_value = ops.multiply(loss_value, loss_weight)
-            return loss_value
+            # Although we are in the fast path, we still need to iterate
+            # through the losses to prevent the torch compiler from failing.
+            loss_values = []
+            for path, loss_fn, loss_weight, _ in self._flat_losses:
+                y_t, y_p = (
+                    resolve_path(path, y_true),
+                    resolve_path(path, y_pred),
+                )
+                if sample_weight is not None and tree.is_nested(sample_weight):
+                    _sample_weight = resolve_path(path, sample_weight)
+                else:
+                    _sample_weight = sample_weight
+                value = ops.cast(
+                    loss_fn(y_t, y_p, _sample_weight), dtype=self.dtype
+                )
+                if loss_weight is not None:
+                    value = ops.multiply(value, loss_weight)
+                loss_values.append(value)
+            return loss_values[0]
 
         try:
             tree.assert_same_structure(y_pred, y_true)
@@ -779,11 +805,6 @@ class CompileLoss(losses_module.Loss):
         # Iterate all losses in flat form.
         loss_values = []
 
-        def resolve_path(path, object):
-            for _path in path:
-                object = object[_path]
-            return object
-
         for (path, loss_fn, loss_weight, _), metric in zip(
             self._flat_losses, metrics
         ):

keras/src/trainers/data_adapters/grain_dataset_adapter.py

@@ -5,13 +5,9 @@ import numpy as np
 from keras.src import tree
 from keras.src.trainers.data_adapters import data_adapter_utils
 from keras.src.trainers.data_adapters.data_adapter import DataAdapter
+from keras.src.utils.module_utils import grain
 from keras.src.utils.module_utils import tensorflow as tf
 
-try:
-    import grain
-except ImportError:
-    grain = None
-
 
 class GrainDatasetAdapter(DataAdapter):
     """Adapter that handles `grain.DataLoader`, `grain.MapDataset` and

keras/src/tree/torchtree_impl.py

@@ -0,0 +1,215 @@
+from collections import defaultdict
+
+from torch.utils import _pytree as torch_tree
+
+
+def register_tree_node_class(cls):
+    torch_tree.register_pytree_node(
+        cls,
+        flatten_fn=lambda x: x.torchtree_flatten(),
+        unflatten_fn=cls.torchtree_unflatten,
+        serialized_type_name=f"{cls.__name__}",
+        flatten_with_keys_fn=lambda x: x.torchtree_flatten_with_keys(),
+    )
+    return cls
+
+
+def _tree_is_leaf(tree, is_leaf=None):
+    if is_leaf is not None and is_leaf(tree):
+        return True
+    return torch_tree._get_node_type(tree) not in torch_tree.SUPPORTED_NODES
+
+
+def _dict_to_ordered_dict(structure):
+    # We need to sort dict and defaultdict to ensure a deterministic order that
+    # that is consistent with other tree implementations.
+    def func(x):
+        if type(x) is dict:
+            return {k: x[k] for k in sorted(x.keys())}
+        elif type(x) is defaultdict:
+            return defaultdict(
+                x.default_factory,
+                {k: x[k] for k in sorted(x.keys())},
+            )
+        return None
+
+    def traverse_children():
+        children, treedef = torch_tree.tree_flatten(
+            structure,
+            is_leaf=lambda x: x is not structure,
+        )
+        if treedef.num_nodes == 1 and treedef.num_leaves == 1:
+            return structure
+        else:
+            return torch_tree.tree_unflatten(
+                [_dict_to_ordered_dict(c) for c in children],
+                treedef,
+            )
+
+    ret = func(structure)
+    if ret is None:
+        return traverse_children()
+    if isinstance(ret, type) and ret.__name__ == "MAP_TO_NONE":
+        return None
+    return ret
+
+
+def is_nested(structure):
+    return not _tree_is_leaf(structure)
+
+
+def traverse(func, structure, top_down=True):
+    def traverse_children():
+        children, treedef = torch_tree.tree_flatten(
+            structure,
+            is_leaf=lambda x: x is not structure,
+        )
+        if treedef.num_nodes == 1 and treedef.num_leaves == 1:
+            return structure
+        else:
+            return torch_tree.tree_unflatten(
+                [traverse(func, c, top_down=top_down) for c in children],
+                treedef,
+            )
+
+    structure = _dict_to_ordered_dict(structure)
+    if top_down:
+        ret = func(structure)
+        if ret is None:
+            return traverse_children()
+    else:
+        traversed_structure = traverse_children()
+        ret = func(traversed_structure)
+        if ret is None:
+            return traversed_structure
+    # Detect MAP_TO_NONE without tree_api import to avoid circular import.
+    if isinstance(ret, type) and ret.__name__ == "MAP_TO_NONE":
+        return None
+    return ret
+
+
+def flatten(structure):
+    # We need to first sort dicts to ensure a deterministic order that is
+    # consistent with other tree implementations.
+    structure = _dict_to_ordered_dict(structure)
+    leaves, _ = torch_tree.tree_flatten(structure)
+    return leaves
+
+
+def flatten_with_path(structure):
+    # We need to first sort dicts to ensure a deterministic order that is
+    # consistent with other tree implementations.
+    structure = _dict_to_ordered_dict(structure)
+    leaves_with_path, _ = torch_tree.tree_flatten_with_path(structure)
+    results = []
+    fields = []
+    for key, leaf in leaves_with_path:
+        for k in key:
+            if isinstance(k, torch_tree.GetAttrKey) and k.name not in fields:
+                fields.append(k.name)
+    fields = sorted(fields)
+    field_to_idx = {f: i for i, f in enumerate(fields)}
+    for key, leaf in leaves_with_path:
+        # Convert to a tuple of keys.
+        path = []
+        for k in key:
+            if isinstance(k, torch_tree.SequenceKey):
+                path.append(k.idx)
+            elif isinstance(k, torch_tree.MappingKey):
+                path.append(k.key)
+            elif isinstance(k, torch_tree.GetAttrKey):
+                path.append(field_to_idx[k.name])
+        results.append((tuple(path), leaf))
+    return results
+
+
+def map_structure(func, *structures, none_is_leaf=True):
+    if not structures:
+        raise ValueError("Must provide at least one structure")
+
+    map_func = func
+    if not none_is_leaf:
+
+        def func_skipping_none(*args):
+            # Check if the reference entry (first one) is None
+            if args[0] is None:
+                if not all(s is None for s in args):
+                    raise ValueError(
+                        "Structure mismatch: some arguments are None, others "
+                        f"are not. Received arguments: {args}."
+                    )
+                return None
+            return func(*args)
+
+        map_func = func_skipping_none
+
+    return torch_tree.tree_map(map_func, *structures)
+
+
+def map_structure_up_to(shallow_structure, func, *structures):
+    if not structures:
+        raise ValueError("Must provide at least one structure")
+
+    # Add check that `shallow_structure` really is the shallowest.
+    # Also only call `func` on `structures` and not `shallow_structure`.
+    def func_with_check_without_shallow_structure(shallow, *args):
+        if not _tree_is_leaf(shallow):
+            raise ValueError("Structures don't have the same nested structure.")
+        return func(*args)
+
+    return torch_tree.tree_map(
+        func_with_check_without_shallow_structure,
+        shallow_structure,
+        *structures,
+    )
+
+
+def assert_same_structure(a, b):
+    def check(a_leaf, b_leaf):
+        if not _tree_is_leaf(a_leaf) or not _tree_is_leaf(b_leaf):
+            raise ValueError("Structures don't have the same nested structure.")
+        return None
+
+    torch_tree.tree_map(check, a, b)
+
+
+def assert_same_paths(a, b):
+    a_paths = set([path for path, _ in flatten_with_path(a)])
+    b_paths = set([path for path, _ in flatten_with_path(b)])
+
+    if a_paths != b_paths:
+        msg = "`a` and `b` don't have the same paths."
+        a_diff = a_paths.difference(b_paths)
+        if a_diff:
+            msg += f"\nPaths in `a` missing in `b`:\n{a_diff}"
+        b_diff = b_paths.difference(a_paths)
+        if b_diff:
+            msg += f"\nPaths in `b` missing in `a`:\n{b_diff}"
+        raise ValueError(msg)
+
+
+def pack_sequence_as(structure, flat_sequence):
+    # We need to first sort dicts to ensure a deterministic order that is
+    # consistent with other tree implementations.
+    structure = _dict_to_ordered_dict(structure)
+    _, treespec = torch_tree.tree_flatten(structure)
+    return torch_tree.tree_unflatten(flat_sequence, treespec)
+
+
+def lists_to_tuples(structure):
+    def list_to_tuple(instance):
+        return tuple(instance) if isinstance(instance, list) else None
+
+    return traverse(list_to_tuple, structure, top_down=False)
+
+
+def map_shape_structure(func, structure):
+    def is_shape_tuple(x):
+        return isinstance(x, (list, tuple)) and all(
+            isinstance(e, (int, type(None))) for e in x
+        )
+
+    # We need to first sort dicts to ensure a deterministic order that is
+    # consistent with other tree implementations.
+    structure = _dict_to_ordered_dict(structure)
+    return torch_tree.tree_map(func, structure, is_leaf=is_shape_tuple)