keras-nightly 3.14.0.dev2025122704__py3-none-any.whl → 3.14.0.dev2026012204__py3-none-any.whl

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/jax_layer.py

@@ -11,6 +11,7 @@ from keras.src.api_export import keras_export
 from keras.src.backend.common.variables import is_float_dtype
 from keras.src.backend.common.variables import standardize_dtype
 from keras.src.layers.layer import Layer
+from keras.src.random.seed_generator import draw_seed
 from keras.src.saving import serialization_lib
 from keras.src.utils import jax_utils
 from keras.src.utils import tracking
@@ -244,15 +245,9 @@ class JaxLayer(Layer):
                 f" Tensorflow backend. Current backend: {backend.backend()}"
             )
 
-        if init_fn is None and params is None and state is None:
-            raise ValueError(
-                "`init_fn`, `params` and `state` cannot all be `None`."
-            )
-
         super().__init__(**kwargs)
         self.call_fn = call_fn
         self.init_fn = init_fn
-        self.seed_generator = backend.random.SeedGenerator(seed)
         self.tracked_params = self._create_variables(params, trainable=True)
         self.tracked_state = self._create_variables(state, trainable=False)
         if self.params is not None or self.state is not None:
@@ -264,7 +259,25 @@ class JaxLayer(Layer):
             {"params", "state", "rng", "inputs", "training"},
             {"inputs"},
         )
-        self.has_state = "state" in self.call_fn_arguments
+        self.call_fn_has_params = "params" in self.call_fn_arguments
+        self.call_fn_has_state = "state" in self.call_fn_arguments
+        call_fn_has_rng = "rng" in self.call_fn_arguments
+
+        if call_fn_has_rng:
+            self.seed_generator = backend.random.SeedGenerator(seed)
+        else:
+            self.seed_generator = None
+
+        if (
+            init_fn is None
+            and params is None
+            and state is None
+            and (self.call_fn_has_params or self.call_fn_has_state)
+        ):
+            raise ValueError(
+                "`init_fn`, `params` and `state` cannot all be `None` when "
+                "`call_fn` takes a `params` or a `state` argument."
+            )
 
         if init_fn:
             self.init_fn_arguments = self._validate_signature(
@@ -428,37 +441,58 @@ class JaxLayer(Layer):
         flat_variables, _ = jax.tree_util.tree_flatten(variables)
         return flat_variables
 
+    def _get_init_seed(self):
+        """
+        Returns a single seed as a tensor of shape [2].
+
+        Call this within `_get_init_rng()` to obtain a new seed.
+
+        Returns:
+            A native tensor of shape [2] and the backend dtype for seeds.
+        """
+        # Use the global SeedGenerator.
+        return draw_seed(None)
+
     def _get_init_rng(self):
         """
-        Returns a key in form of the backend array of size 2 dtype uint32
-        to pass to `init_fn`.
+        Returns a seed or seeds to pass as the `rng` argument of `init_fn`.
+
+        By default, this returns a single seed. Override this to return a
+        different structure. Overrides should use `self._get_init_seed()` to
+        obtain new seeds.
+
+        Returns:
+            RNG key or structure of keys as tensors of shape [2] and the backend
+            dtype for seeds.
+        """
+        return self._get_init_seed()
+
+    def _get_call_seed(self):
+        """
+        Returns a single seed as a tensor of shape [2].
 
-        By default, this returns a Jax or TF array of size 2 by calling
-        `self.seed_generator.next()`. Override this to return a different
-        structure.
+        Call this within `_get_call_rng()` to obtain a new seed.
 
         Returns:
-            a key as an Jax or TF array of size 2 dtype uint32 will be passed
-            as the `rng` argument of `init_fn`.
+            A native tensor of shape [2] and the backend dtype for seeds.
         """
         return self.seed_generator.next()
 
     def _get_call_rng(self, training):
         """
-        Returns a key in form of the backend array of size 2 dtype uint32
-        to pass to `call_fn`.
+        Returns a seed or seeds to pass as the `rng` argument of `call_fn`.
 
-        By default, this returns a Jax or TF array of size 2 by calling
-        `self.seed_generator.next()` when `training` is `True`, and `None` when
-        `training` is `False`. Override this to return a different structure or
-        to pass RNGs in inference mode too.
+        By default, this returns a seed when `training` is `True`, and `None`
+        when `training` is `False`. Override this to return a different
+        structure or to pass seeds in inference mode too. Overrides should use
+        `self._get_call_seed()` to obtain seeds.
 
         Returns:
-            a key as an Jax or TF array of size 2 dtype uint32 will be passed
-            as the `rng` argument of `call_fn`.
+            RNG key or structure of keys as tensors of shape [2] and the backend
+            dtype for seeds.
         """
         if training:
-            return self.seed_generator.next()
+            return self._get_call_seed()
         else:
             return None
 
@@ -492,7 +526,7 @@ class JaxLayer(Layer):
                 init_args.append(True)
 
         init_result = self.init_fn(*init_args)
-        if self.has_state:
+        if self.call_fn_has_state:
             init_params, init_state = init_result
         else:
             init_params, init_state = init_result, None
@@ -503,7 +537,11 @@ class JaxLayer(Layer):
         self.tracked_state = self._create_variables(init_state, trainable=False)
 
     def build(self, input_shape):
-        if self.params is None and self.state is None:
+        if (
+            self.params is None
+            and self.state is None
+            and (self.call_fn_has_params or self.call_fn_has_state)
+        ):
             self._initialize_weights(input_shape)
 
         if backend.backend() == "tensorflow":
@@ -578,7 +616,7 @@ class JaxLayer(Layer):
             variable.assign(value)
 
         def call_with_fn(fn):
-            if self.has_state:
+            if self.call_fn_has_state:
                 predictions, new_state = fn(*call_args)
                 jax.tree_util.tree_map(
                     assign_state_to_variable, new_state, self.state
@@ -711,12 +749,12 @@ class FlaxLayer(JaxLayer):
         **kwargs,
     ):
         # Late import to only require Flax when this is used.
-        from flax.core import scope as flax_scope
+        from flax.linen import DenyList
 
         self.module = module
         self.method = method
 
-        apply_mutable = flax_scope.DenyList(["params"])
+        apply_mutable = DenyList(["params"])
 
         def apply_with_training(params, state, rng, inputs, training):
             return self.module.apply(
@@ -801,13 +839,13 @@ class FlaxLayer(JaxLayer):
 
     def _get_init_rng(self):
         return {
-            "params": self.seed_generator.next(),
-            "dropout": self.seed_generator.next(),
+            "params": self._get_init_seed(),
+            "dropout": self._get_init_seed(),
         }
 
     def _get_call_rng(self, training):
         if training:
-            return {"dropout": self.seed_generator.next()}
+            return {"dropout": self._get_call_seed()}
         else:
             return {}
 

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/utils/tracking.py

@@ -31,13 +31,13 @@ def no_automatic_dependency_tracking(fn):
 class Tracker:
     """Attribute tracker, used for e.g. Variable tracking.
 
-    Monitors certain attribute types
-    and put them in appropriate lists in case of a match.
+    Monitors certain attribute types and places matching
+    objects into user provided tracking collections.
 
     Also passively tracks certain mutable collections
-    (dict, list) so that items added to them later
-    still get tracked. This is done by wrapping these
-    collections into an equivalent, tracking-aware object.
+    (e.g. dict and list) ensuring that items added after
+    initialization are still tracked. This is done by wrapping
+    these collections in tracking-aware proxy objects.
 
     Example:
 

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.dev2025122704"
+__version__ = "3.14.0.dev2026012204"
 
 
 @keras_export("keras.version")

{keras_nightly-3.14.0.dev2025122704.dist-info → keras_nightly-3.14.0.dev2026012204.dist-info}/METADATA

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