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

keras/src/utils/jax_layer.py

@@ -1,4 +1,7 @@
+import functools
 import inspect
+import itertools
+import string
 
 import numpy as np
 
@@ -8,10 +11,27 @@ 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
 from keras.src.utils.module_utils import jax
+from keras.src.utils.module_utils import tensorflow as tf
+
+if backend.backend() == "tensorflow":
+    tf_no_automatic_dependency_tracking = (
+        tf.__internal__.tracking.no_automatic_dependency_tracking
+    )
+else:
+
+    def tf_no_automatic_dependency_tracking(fn):
+        return fn
+
+
+def _convert_to_jax_key(tensor):
+    if backend.backend() == "tensorflow":
+        return tf.bitcast(tensor, tf.uint32)[0]
+    return tensor
 
 
 @keras_export("keras.layers.JaxLayer")
@@ -219,21 +239,15 @@ class JaxLayer(Layer):
         seed=None,
         **kwargs,
     ):
-        if backend.backend() != "jax":
+        if backend.backend() not in ["jax", "tensorflow"]:
             raise ValueError(
-                "JaxLayer is only supported with the JAX backend. Current "
-                f"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`."
+                f"{self.__class__.__name__} is only supported with the JAX or"
+                f" Tensorflow backend. Current backend: {backend.backend()}"
             )
 
         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:
@@ -245,13 +259,35 @@ 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(
                 init_fn, "init_fn", {"rng", "inputs", "training"}, {"inputs"}
             )
 
+        # Attributes for jax2tf functions
+        self.jax2tf_training_false_fn = None
+        self.jax2tf_training_true_fn = None
+
     def _validate_signature(self, fn, fn_name, allowed, required):
         fn_parameters = inspect.signature(fn).parameters
         for parameter_name in required:
@@ -272,7 +308,81 @@ class JaxLayer(Layer):
 
         return parameter_names
 
+    def _get_jax2tf_input_shape(self, input_shape):
+        """Convert input shape in a format suitable for `jax2tf`.
+
+        `jax2tf` expects a letter for each unknown dimension, which allows
+        correlated dimensions. Since correlated dimensions are not supported by
+        Keras, we simply use 'a', 'b', 'c'..., for each unknown dimension. We
+        however use 'batch' for dimension 0 if not defined to correlate the
+        batch size across inputs.
+
+        Example (spaces added for readability):
+        ```
+        input_shape:  (None , 4   , None, None, 5   )
+        result:      "(batch, 4   , a   , b   , 5   )"
+        ```
+
+        Args:
+          input_shape: a single shape or a structure of shapes for the inputs.
+        Returns:
+          the shape or shapes structure in the `jax2tf` format as strings.
+        """
+        dim_names = itertools.chain(
+            string.ascii_lowercase,  # a, b, ... z
+            itertools.starmap(  # aa, ab, ... az, ba, bb, ... zz
+                lambda a, b: a + b,
+                itertools.product(string.ascii_lowercase, repeat=2),
+            ),
+        )
+
+        def get_single_jax2tf_shape(shape):
+            jax2tf_shape = []
+
+            for index, dim in enumerate(shape):
+                if dim is not None:
+                    jax2tf_shape.append(str(dim))
+                elif index == 0:
+                    jax2tf_shape.append("batch")
+                else:
+                    jax2tf_shape.append(next(dim_names))
+
+            return "(" + ", ".join(jax2tf_shape) + ")"
+
+        res = tree.map_shape_structure(get_single_jax2tf_shape, input_shape)
+        return res
+
+    def _jax2tf_convert(self, fn, polymorphic_shapes):
+        from jax.experimental import jax2tf
+
+        converted_fn = jax2tf.convert(fn, polymorphic_shapes=polymorphic_shapes)
+        # Autograph won't work with the output of jax2tf.
+        converted_fn = tf.autograph.experimental.do_not_convert(converted_fn)
+        return converted_fn
+
+    def _partial_with_positional(self, fn, index, value):
+        """Return a new partial with one positional argument set to a value.
+
+        This is needed because `jax2tf` only supports positional arguments and
+        `functools.partial` only supports setting positional arguments starting
+        from the left. Our use case is the `training` argument which is
+        typically the righmost argument.
+
+        Args:
+          fn: the function to wrap.
+          index: the index of the positional argument to set to `value`.
+          value: the value for the positional argument at `index`.
+        """
+
+        @functools.wraps(fn)
+        def wrapper(*args):
+            args = args[0:index] + (value,) + args[index:]
+            return fn(*args)
+
+        return wrapper
+
     @tracking.no_automatic_dependency_tracking
+    @tf_no_automatic_dependency_tracking
     def _create_variables(self, values, trainable):
         """Create a structure of variables from a structure of JAX arrays.
 
@@ -296,14 +406,14 @@ class JaxLayer(Layer):
 
         def create_variable(value):
             if backend.is_tensor(value) or isinstance(
-                value, (np.ndarray, np.generic)
+                value, (np.ndarray, np.generic, jax.Array)
             ):
                 dtype = value.dtype
                 if is_float_dtype(dtype):
                     dtype = None  # Use the layer dtype policy
                 return self.add_weight(
                     value.shape,
-                    initializer=value,
+                    initializer=backend.convert_to_tensor(value),
                     dtype=dtype,
                     trainable=trainable,
                 )
@@ -331,46 +441,69 @@ 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 JAX `PRNGKey` or structure of `PRNGKey`s 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 `PRNGKey` retrieved by calling
-        `self.seed_generator.next()`. Override this to return a different
-        structure.
+        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:
-            a JAX `PRNGKey` or structure of `PRNGKey`s that will be passed as
-            the `rng` argument of `init_fn`.
+            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].
+
+        Call this within `_get_call_rng()` to obtain a new seed.
+
+        Returns:
+            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 JAX `PRNGKey` or structure of `PRNGKey`s to pass to `call_fn`.
+        Returns a seed or seeds to pass as the `rng` argument of `call_fn`.
 
-        By default, this returns a single `PRNGKey` retrieved 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 JAX `PRNGKey` or structure of `PRNGKey`s that 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
 
-    def build(self, input_shape):
-        if self.params is not None or self.state is not None:
-            return
-
-        if jax_utils.is_in_jax_tracing_scope():
+    def _initialize_weights(self, input_shape):
+        if jax_utils.is_in_jax_tracing_scope() or tf.inside_function():
             # This exception is not actually shown, it is caught and a detailed
             # warning about calling 'build' is printed.
-            raise ValueError("'JaxLayer' cannot be built in tracing scope")
+            raise ValueError(
+                "'JaxLayer' cannot be built in tracing scope"
+                "or inside tf function"
+            )
 
         # Initialize `params` and `state` if needed by calling `init_fn`.
         def create_input(shape):
@@ -381,14 +514,19 @@ class JaxLayer(Layer):
         init_args = []
         for argument_name in self.init_fn_arguments:
             if argument_name == "rng":
-                init_args.append(self._get_init_rng())
+                init_args.append(
+                    jax.tree_util.tree_map(
+                        lambda x: jax.numpy.array(_convert_to_jax_key(x)),
+                        self._get_init_rng(),
+                    )
+                )
             elif argument_name == "inputs":
                 init_args.append(init_inputs)
             elif argument_name == "training":
                 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
@@ -398,6 +536,49 @@ 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
+            and (self.call_fn_has_params or self.call_fn_has_state)
+        ):
+            self._initialize_weights(input_shape)
+
+        if backend.backend() == "tensorflow":
+            polymorphic_shapes = []
+            for argument in self.call_fn_arguments:
+                if argument == "inputs":
+                    polymorphic_shapes.append(
+                        self._get_jax2tf_input_shape(input_shape)
+                    )
+                elif argument != "training":
+                    # params, state, rng
+                    polymorphic_shapes.append("...")
+
+            if "training" in self.call_fn_arguments:
+                training_argument_index = self.call_fn_arguments.index(
+                    "training"
+                )
+                self.jax2tf_training_false_fn = self._jax2tf_convert(
+                    self._partial_with_positional(
+                        self.call_fn, training_argument_index, False
+                    ),
+                    polymorphic_shapes,
+                )
+                self.jax2tf_training_true_fn = self._jax2tf_convert(
+                    self._partial_with_positional(
+                        self.call_fn, training_argument_index, True
+                    ),
+                    polymorphic_shapes,
+                )
+            else:
+                self.jax2tf_training_false_fn = self._jax2tf_convert(
+                    self.call_fn,
+                    polymorphic_shapes,
+                )
+                self.jax2tf_training_true_fn = None
+            super().build(input_shape)
+
     def call(self, inputs, training=False):
         def unwrap_variable(variable):
             return None if variable is None else variable.value
@@ -413,11 +594,16 @@ class JaxLayer(Layer):
                     jax.tree_util.tree_map(unwrap_variable, self.state)
                 )
             elif argument_name == "rng":
-                call_args.append(self._get_call_rng(training))
+                call_args.append(
+                    jax.tree_util.tree_map(
+                        _convert_to_jax_key, self._get_call_rng(training)
+                    )
+                )
             elif argument_name == "inputs":
                 call_args.append(inputs)
             elif argument_name == "training":
-                call_args.append(training)
+                if backend.backend() == "jax":
+                    call_args.append(training)
 
         def assign_state_to_variable(value, variable):
             # This exists only to make debugging this error case easier.
@@ -429,14 +615,23 @@ class JaxLayer(Layer):
                 )
             variable.assign(value)
 
-        if self.has_state:
-            predictions, new_state = self.call_fn(*call_args)
-            jax.tree_util.tree_map(
-                assign_state_to_variable, new_state, self.state
-            )
-            return predictions
-        else:
-            return self.call_fn(*call_args)
+        def call_with_fn(fn):
+            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
+                )
+                return predictions
+            else:
+                return fn(*call_args)
+
+        if backend.backend() == "jax":
+            return call_with_fn(self.call_fn)
+        elif backend.backend() == "tensorflow":
+            if training and self.jax2tf_training_true_fn is not None:
+                return call_with_fn(self.jax2tf_training_true_fn)
+            else:
+                return call_with_fn(self.jax2tf_training_false_fn)
 
     def get_config(self):
         config = {
@@ -554,18 +749,12 @@ class FlaxLayer(JaxLayer):
         **kwargs,
     ):
         # Late import to only require Flax when this is used.
-        from flax.core import scope as flax_scope
-
-        if backend.backend() != "jax":
-            raise ValueError(
-                "FlaxLayer is only supported with the JAX backend. Current "
-                f"backend: {backend.backend()}"
-            )
+        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(
@@ -650,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

@@ -39,11 +39,31 @@ class LazyModule:
         return f"LazyModule({self.name})"
 
 
+class OrbaxLazyModule(LazyModule):
+    def initialize(self):
+        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=(
@@ -59,3 +79,12 @@ optree = LazyModule("optree")
 dmtree = LazyModule("tree")
 tf2onnx = LazyModule("tf2onnx")
 grain = LazyModule("grain")
+litert = LazyModule("ai_edge_litert")
+ocp = OrbaxLazyModule(
+    "orbax.checkpoint.v1",
+    pip_name="orbax-checkpoint",
+    import_error_msg=(
+        "OrbaxCheckpoint requires the 'orbax-checkpoint' package. "
+        "You can install it via pip install orbax-checkpoint"
+    ),
+)

keras/src/utils/progbar.py

@@ -3,7 +3,8 @@ import os
 import sys
 import time
 
-from keras.src import backend
+import numpy as np
+
 from keras.src.api_export import keras_export
 from keras.src.utils import io_utils
 
@@ -162,12 +163,10 @@ class Progbar:
             for k in self._values_order:
                 info += f" - {k}:"
                 if isinstance(self._values[k], list):
-                    avg = backend.convert_to_numpy(
-                        backend.numpy.mean(
-                            self._values[k][0] / max(1, self._values[k][1])
-                        )
-                    )
-                    avg = float(avg)
+                    values, count = self._values[k]
+                    if not isinstance(values, float):
+                        values = np.mean(values)
+                    avg = values / max(1, count)
                     if abs(avg) > 1e-3:
                         info += f" {avg:.4f}"
                     else:
@@ -194,11 +193,10 @@ class Progbar:
                 info += f" -{self._format_time(time_per_unit, self.unit_name)}"
                 for k in self._values_order:
                     info += f" - {k}:"
-                    avg = backend.convert_to_numpy(
-                        backend.numpy.mean(
-                            self._values[k][0] / max(1, self._values[k][1])
-                        )
-                    )
+                    values, count = self._values[k]
+                    if not isinstance(values, float):
+                        values = np.mean(values)
+                    avg = values / max(1, count)
                     if avg > 1e-3:
                         info += f" {avg:.4f}"
                     else:

keras/src/utils/python_utils.py

@@ -181,6 +181,8 @@ def pythonify_logs(logs):
         A flattened dict with values converted to Python-native types if
         possible.
     """
+    from keras.src import backend
+
     logs = logs or {}
     result = {}
     for key, value in sorted(logs.items()):
@@ -188,6 +190,9 @@ def pythonify_logs(logs):
             result.update(pythonify_logs(value))
         else:
             try:
+                # Prevent torch compiler from breaking the graph.
+                if backend.is_tensor(value):
+                    value = backend.convert_to_numpy(value)
                 value = float(value)
             except:
                 pass

keras/src/utils/rng_utils.py

@@ -5,6 +5,7 @@ import numpy as np
 from keras.src import backend
 from keras.src.api_export import keras_export
 from keras.src.backend.common import global_state
+from keras.src.random import seed_generator
 from keras.src.utils.module_utils import tensorflow as tf
 
 GLOBAL_RANDOM_SEED = "global_random_seed"
@@ -20,7 +21,7 @@ def set_random_seed(seed):
     sources of randomness, or when certain non-deterministic cuDNN ops are
     involved.
 
-    Calling this utility is equivalent to the following:
+    Calling this utility does the following:
 
     ```python
     import random
@@ -36,6 +37,9 @@ def set_random_seed(seed):
     torch.manual_seed(seed)
     ```
 
+    Additionally, it resets the global Keras `SeedGenerator`, which is used by
+    `keras.random` functions when the `seed` is not provided.
+
     Note that the TensorFlow seed is set even if you're not using TensorFlow
     as your backend framework, since many workflows leverage `tf.data`
     pipelines (which feature random shuffling). Likewise many workflows
@@ -52,6 +56,10 @@ def set_random_seed(seed):
 
     # Store seed in global state so we can query it if set.
     global_state.set_global_attribute(GLOBAL_RANDOM_SEED, seed)
+    # Remove global SeedGenerator, it will be recreated from the seed.
+    global_state.set_global_attribute(
+        seed_generator.GLOBAL_SEED_GENERATOR, None
+    )
     random.seed(seed)
     np.random.seed(seed)
     if tf.available: