keras-nightly 3.14.0.dev2026011504__py3-none-any.whl → 3.14.0.dev2026011704__py3-none-any.whl

keras/src/backend/tensorflow/rnn.py

@@ -539,11 +539,21 @@ def _do_lstm_arguments_support_cudnn(
 
 
 def _has_fully_masked_sequence(mask):
-    # Cudnn kernel will error out if the input sequence contains any
-    # fully masked data. We walk around this issue by rerouting the computation
-    # to standard kernel, until the issue on cudnn side has been fixed.  For a
-    # fully masked sequence, it will contain all Falses. To make it easy to
-    # check, we inverse the boolean, check if any of the sequence has all True.
+    """Check if input sequence contains any fully masked data.
+
+    cuDNN kernel will error out if the input sequence contains any fully masked
+    data. We work around this issue by rerouting the computation to the
+    standard kernel until the issue on the cuDNN side has been fixed. For a
+    fully masked sequence, it will contain all `False` values. To make it easy
+    to check, we invert the boolean and check if any of the sequences has all
+    `True` values.
+
+    Args:
+        mask: The mask tensor.
+
+    Returns:
+        A boolean tensor, `True` if the mask contains a fully masked sequence.
+    """
     return tf.reduce_any(
         tf.reduce_all(tf.logical_not(tf.cast(mask, dtype="bool")), axis=1)
     )
@@ -900,8 +910,8 @@ def _cudnn_lstm(
 
     if tf.sysconfig.get_build_info()["is_rocm_build"]:
         # ROCm MIOpen's weight sequence for LSTM is different from both
-        # canonical and Cudnn format
-        # MIOpen: [i, f, o, c] Cudnn/Canonical: [i, f, c, o]
+        # canonical and cuDNN format
+        # MIOpen: [i, f, o, c] cuDNN/Canonical: [i, f, c, o]
         # i is input gate weights.
         # f is forget gate weights.
         # o is output gate weights.

keras/src/backend/torch/rnn.py

@@ -413,11 +413,21 @@ def _is_sequence_right_padded(mask):
 
 
 def _has_fully_masked_sequence(mask):
-    # Cudnn kernel will error out if the input sequence contains any
-    # fully masked data. We walk around this issue by rerouting the computation
-    # to standard kernel, until the issue on cudnn side has been fixed.  For a
-    # fully masked sequence, it will contain all Falses. To make it easy to
-    # check, we inverse the boolean, check if any of the sequence has all True.
+    """Check if input sequence contains any fully masked data.
+
+    cuDNN kernel will error out if the input sequence contains any fully masked
+    data. We work around this issue by rerouting the computation to the
+    standard kernel until the issue on the cuDNN side has been fixed. For a
+    fully masked sequence, it will contain all `False` values. To make it easy
+    to check, we invert the boolean and check if any of the sequences has all
+    `True` values.
+
+    Args:
+        mask: The mask tensor.
+
+    Returns:
+        A boolean tensor, `True` if the mask contains a fully masked sequence.
+    """
     return torch.any(torch.all(~mask, dim=1))
 
 
@@ -447,8 +457,8 @@ def _compute_sequence_length_from_mask(mask, batch_first):
     The masking tensor is a 2D boolean tensor with shape [batch, timestep]. For
     any timestep that should be masked, the corresponding field will be False.
     Consider the following example:
-      a = [[True, True, False, False]
-           [True, True, True, False]]
+        a = [[True, True, False, False]
+             [True, True, True, False]]
     It is a (2, 4) tensor, and the corresponding sequence length result should
     be 1D tensor with value [2, 3]. Note that the masking tensor must be right
     padded that could be checked by, e.g., `is_sequence_right_padded()`.
@@ -467,12 +477,19 @@ def _compute_sequence_length_from_mask(mask, batch_first):
 
 
 def prepare_lstm_weights(lstm, kernel, recurrent_kernel, bias, device):
-    """Copies kernel and recurrent kernel weights in the Pytorch format
+    """Copies kernel and recurrent kernel weights into the PyTorch format.
+
     We split the kernel and recurrent kernel weights, create associated
-    torch tensors adapted to be in line with the Cudnn optimization.
-    After we have copied the weights, we ensure the paramters are on
-    the same device and memory layout is optimized for Cudnn.
+    torch tensors adapted to be in line with the cuDNN optimization.
+    After we have copied the weights, we ensure the parameters are on
+    the same device and memory layout is optimized for cuDNN.
 
+    Args:
+        lstm: The PyTorch LSTM layer to prepare weights for.
+        kernel: The kernel weights tensor.
+        recurrent_kernel: The recurrent kernel weights tensor.
+        bias: The bias tensor.
+        device: The device to place the tensors on.
     """
 
     lstm = lstm.to(device)

keras/src/layers/preprocessing/image_preprocessing/random_contrast.py

@@ -92,8 +92,8 @@ class RandomContrast(BaseImagePreprocessingLayer):
 
     def transform_images(self, images, transformation, training=True):
         if training:
-            constrast_factor = transformation["contrast_factor"]
-            outputs = self._adjust_constrast(images, constrast_factor)
+            contrast_factor = transformation["contrast_factor"]
+            outputs = self._adjust_contrast(images, contrast_factor)
             outputs = self.backend.numpy.clip(
                 outputs, self.value_range[0], self.value_range[1]
             )
@@ -117,7 +117,7 @@ class RandomContrast(BaseImagePreprocessingLayer):
     ):
         return segmentation_masks
 
-    def _adjust_constrast(self, inputs, contrast_factor):
+    def _adjust_contrast(self, inputs, contrast_factor):
         if self.data_format == "channels_first":
             height_axis = -2
             width_axis = -1

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

{keras_nightly-3.14.0.dev2026011504.dist-info → keras_nightly-3.14.0.dev2026011704.dist-info}/METADATA

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

{keras_nightly-3.14.0.dev2026011504.dist-info → keras_nightly-3.14.0.dev2026011704.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=9qMnmtF-qZTXfhu_aYy3T9wAf18rbTFfpRqEiny8QSU,204
+keras/src/version.py,sha256=PYl1X5NcUeyqMJlZOu02EORfq3XehPpN28bC457e3F8,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
@@ -214,7 +214,7 @@ keras/src/backend/tensorflow/nn.py,sha256=6vtZHzUED6_blUPE1Tnc3GAxPpJ2ebxoaiMn80
 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
+keras/src/backend/tensorflow/rnn.py,sha256=JbOSpt48cm612c7YwiTYOQCQsNXyI_6QeRhtUn8qEvM,34829
 keras/src/backend/tensorflow/sparse.py,sha256=a_FZcJY-wPl1x4vY0T7j-GORa4SAuMjNEToJLmK0daQ,32247
 keras/src/backend/tensorflow/tensorboard.py,sha256=e7pXicuMfQjuCmq1wOmixWhWt2EbjLMBo_JPAqCbZRk,504
 keras/src/backend/tensorflow/trackable.py,sha256=QZn0JvpBJ7Kx4e6zM2IVIWz9ADcWDB-dHN6vjoQBa9Q,1993
@@ -229,7 +229,7 @@ keras/src/backend/torch/math.py,sha256=g-ElDii2Y_o1-t6BAu2nbS7JH-aPqVS5Fqds8aYzI
 keras/src/backend/torch/nn.py,sha256=zmEzXEuwD7fVRDm145zsxzUDmqNmRgZS4LmeIx4Nbus,37498
 keras/src/backend/torch/numpy.py,sha256=gvHviedkAoEaTax89wDqUrjbUSX1ndjxicHy-PLv2Nc,57668
 keras/src/backend/torch/random.py,sha256=YhLfC7qkGpzlU_i6gGPVormo3BMSo7OUA3TC3GCehrA,8292
-keras/src/backend/torch/rnn.py,sha256=J0vg7ikxBiv1FzEavgwT8IVCs0ceBcEv5LYyM5C2suA,25545
+keras/src/backend/torch/rnn.py,sha256=MJIVbHKsUA2dZm4Gu2NvRxlrFCWeWSxSZRmFxSsC3Zg,26041
 keras/src/backend/torch/trainer.py,sha256=dcikz1c5O0FHNzRKSi6WhIHsHfLV2HDlrXPElSd1cgE,17985
 keras/src/backend/torch/optimizers/__init__.py,sha256=yvqiyKgMEh-nGpacssdpsMySujyYB6lPy-Wil3onXvo,78
 keras/src/backend/torch/optimizers/torch_adadelta.py,sha256=iPjGHvD7q_VD0WaMNxuNcvz8uIWd0smRyEMzMqryUD4,1672
@@ -401,7 +401,7 @@ keras/src/layers/preprocessing/image_preprocessing/rand_augment.py,sha256=upDdEg
 keras/src/layers/preprocessing/image_preprocessing/random_brightness.py,sha256=Ix01T1xsbf_QknyWcSlK1SxVPvFNtHw20xmWHhuQPZI,6083
 keras/src/layers/preprocessing/image_preprocessing/random_color_degeneration.py,sha256=N6rCXPhWCEh-xWqC9ETYwrbJ2f6lIqyCR9Z18uV3xd0,4896
 keras/src/layers/preprocessing/image_preprocessing/random_color_jitter.py,sha256=rbQvLhCPPXyAaYfcMiVzyN0yvfFrcfbRbkVruO9o38U,9464
-keras/src/layers/preprocessing/image_preprocessing/random_contrast.py,sha256=GvB5iQngY-4v99mGS9dXOlGTX4GB6Z7ZvDqW1TKJR5A,5474
+keras/src/layers/preprocessing/image_preprocessing/random_contrast.py,sha256=eJ7aakES1YfSv1JXjv8ZT3ltTqgG6Oo1_XU6BopKDng,5470
 keras/src/layers/preprocessing/image_preprocessing/random_crop.py,sha256=y2iHw-xbSV11uK4D34VT9QEkpvKOk-D-TmVSCZUjDn0,10553
 keras/src/layers/preprocessing/image_preprocessing/random_elastic_transform.py,sha256=fIfPe-906LUhTUDpiuPwM5oEOJ_1UQ9BhMHBFpItcGM,10208
 keras/src/layers/preprocessing/image_preprocessing/random_erasing.py,sha256=O7f44V805Wta9RMZyks4sl-LViglTCdp7_n-qj_nWbI,11233
@@ -585,7 +585,7 @@ keras/src/utils/grain_utils.py,sha256=Wfwv12E3UrNZjJjTEk2JVV6_YEUav35UJ6bV1UAPEI
 keras/src/utils/image_dataset_utils.py,sha256=0lOzD1CiXwZOe1wW-5uvFKuIgot9PWUC9KJJA0NVuP8,24017
 keras/src/utils/image_utils.py,sha256=lGe4iKYQkQ6j15CbHoqpSMC6JEvCrekYBuYGoMClcpo,17051
 keras/src/utils/io_utils.py,sha256=Riv9TCCnz6xQLUvR1QC-UOCoGZ_KiNTwQVvLY6dKcX8,4432
-keras/src/utils/jax_layer.py,sha256=ytws8NcxWzJ4kViBy3bc-Pk3st3_3L8RqXxgq9sYp1k,32912
+keras/src/utils/jax_layer.py,sha256=xwUkk-yp5lieC_uJesn4T4Lkw1bdjtSY5Q-bK8PuHH0,34027
 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=FTZPMRLurURchLPX1tu-h3b-UoPW28faNOlDzpYDW6A,2894
@@ -618,7 +618,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.dev2026011504.dist-info/METADATA,sha256=6s-lhD6ZQgn-dl1KiHGJMLNttJ4ir79glEAKZAMn6fI,6339
-keras_nightly-3.14.0.dev2026011504.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_nightly-3.14.0.dev2026011504.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
-keras_nightly-3.14.0.dev2026011504.dist-info/RECORD,,
+keras_nightly-3.14.0.dev2026011704.dist-info/METADATA,sha256=XtouV2KcEzUqH0W897TEDF7jmTiNPzWoJyzib0rfKAo,6339
+keras_nightly-3.14.0.dev2026011704.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_nightly-3.14.0.dev2026011704.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
+keras_nightly-3.14.0.dev2026011704.dist-info/RECORD,,