keras-nightly 3.14.0.dev2026010204__py3-none-any.whl → 3.14.0.dev2026010304__py3-none-any.whl

keras/src/backend/jax/nn.py

@@ -1471,25 +1471,42 @@ def _can_use_flash_attention(query, key, value, bias, raise_error=False):
         # Only support at least Ampere
         if not check_compute_capability("8.0"):
             raise RuntimeError("Require at least Ampere arch to run")
-        # Check inputs layout
+
+        # Inspect inputs of `check_layout`
         check_layout_params = list(
             inspect.signature(check_layout).parameters.keys()
         )
         for known_param in ("query", "key", "value", "bias", "layout"):
             check_layout_params.remove(known_param)
         # Defaults to `None` when not specified.
-        kwargs = {key: None for key in check_layout_params}
+        check_layout_kwargs = {key: None for key in check_layout_params}
         check_layout(
-            query, key, value, bias, layout=_normalize_layout("BTNH"), **kwargs
-        )
-        check_is_flash_attention(
             query,
             key,
-            _normalize_layout("BTNH"),
-            cudnn_version,
-            bias is not None,
-            is_training=False,
+            value,
+            bias,
+            layout=_normalize_layout("BTNH"),
+            **check_layout_kwargs,
         )
+
+        # Inspect inputs of `check_is_flash_attention`
+        check_is_flash_attention_params = inspect.signature(
+            check_is_flash_attention
+        ).parameters
+        check_is_flash_attention_kwargs = {
+            "query": query,
+            "key": key,
+            "value": value,
+            "layout": _normalize_layout("BTNH"),
+            "cudnn_version": cudnn_version,
+            "has_bias": bias is not None,
+            "is_training": False,
+        }
+        # Remove unsupported arguments
+        for param in list(check_is_flash_attention_kwargs.keys()):
+            if param not in check_is_flash_attention_params:
+                check_is_flash_attention_kwargs.pop(param)
+        check_is_flash_attention(**check_is_flash_attention_kwargs)
         return True
     except:
         if raise_error:

keras/src/random/seed_generator.py

@@ -29,7 +29,7 @@ class SeedGenerator:
     a local `StateGenerator` with either a deterministic or random initial
     state.
 
-    Remark concerning the JAX backen: Note that the use of a local
+    Remark concerning the JAX backend: Note that the use of a local
     `StateGenerator` as seed argument is required for JIT compilation of
     RNG with the JAX backend, because the use of global state is not
     supported.
@@ -111,7 +111,7 @@ class SeedGenerator:
         return new_seed_value
 
     def get_config(self):
-        return {"seed": self._initial_seed}
+        return {"seed": self._initial_seed, "name": self.name}
 
     @classmethod
     def from_config(cls, config):

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

{keras_nightly-3.14.0.dev2026010204.dist-info → keras_nightly-3.14.0.dev2026010304.dist-info}/METADATA

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

{keras_nightly-3.14.0.dev2026010204.dist-info → keras_nightly-3.14.0.dev2026010304.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=qy2X7ZAxdvBc9hNdvShnkuP-3oxQHJY7Wac1BCYm_Hs,204
+keras/src/version.py,sha256=g0zaAy91Gg1cY-Ey6cBCCgMseDoEXzoyCPMqyyGWO-g,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
@@ -170,7 +170,7 @@ keras/src/backend/jax/image.py,sha256=RiYIalbIaUQdDOGpDZUBk5KNsX94Xqg7iyXGATN9V5
 keras/src/backend/jax/layer.py,sha256=o6CicT06udwamTRQIjNSDLZLyYHFzBXNbxewXgWe0iw,308
 keras/src/backend/jax/linalg.py,sha256=LDaLZYz49ChE2kJR3YpaM9xuwusvd3krV7nNAAazTWA,2642
 keras/src/backend/jax/math.py,sha256=1IEDpdoF8e5ltu3D4wbDQuihzvJHhMXz8W9Z_E-eJqU,9391
-keras/src/backend/jax/nn.py,sha256=mQAKxZMedpv6H4GSU_ofmWKpdCbLukJA9Ncx3kfRuVc,59605
+keras/src/backend/jax/nn.py,sha256=mrRawNvf9EWe8rdTwK_Auz6xdLkVG6hH0nIAP7hyUDE,60271
 keras/src/backend/jax/numpy.py,sha256=SMa6dH1n7v04SsnEkevCWBqmzj7Ed8TmBASOSrEQIMM,38619
 keras/src/backend/jax/optimizer.py,sha256=5DeXQHcYmUI6F9i1m1VHn3sBt4LEStOeBXnKdESevLM,4134
 keras/src/backend/jax/random.py,sha256=Uk2huGIk_dlzMrx5eDVrrr2TeCEMitn2vr4yzA0NXjs,3594
@@ -536,7 +536,7 @@ keras/src/quantizers/quantizers.py,sha256=BDD3vi_15lmOY_ybI7oQDgINYlM9CF0QSQuP6k
 keras/src/quantizers/utils.py,sha256=i6e5MobXrQeKA6zFenjzUNoDDWRGL9bcfgdbE_-0IeM,672
 keras/src/random/__init__.py,sha256=BmXVYPzxbhADohoLtAEEzB3cesP7YBFDsp1qc6BWWlg,420
 keras/src/random/random.py,sha256=bUADZIVDuCghwIWTk0qBxXTxUdiNGWIdsRi8QJ3ePg4,17581
-keras/src/random/seed_generator.py,sha256=Iyx_YbLSaYusmCm2rOOOiDNU57x9-sU-xDf8U_YVpTE,5635
+keras/src/random/seed_generator.py,sha256=-a0CQa7--Xt0g0nfdjLmUzlFElY9Y838VcCx05AcllY,5655
 keras/src/regularizers/__init__.py,sha256=GzK9FTKL2Xxd5H55GfG9gxDqt4eZoVHFWICgb2VW8qM,1731
 keras/src/regularizers/regularizers.py,sha256=urXNmMGuqHT7lOmS-yQPl3At3Ny-37Xlo389ErCg84A,11799
 keras/src/saving/__init__.py,sha256=vnrtfvnzW7Gwtxe5COhaMoEnVYB5iDe2YlqJ-DvqFIk,614
@@ -614,7 +614,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.dev2026010204.dist-info/METADATA,sha256=K8uCZBLeOjNnygheaY4ZM3Q076vdcfFxYlFTQpBvgSc,6339
-keras_nightly-3.14.0.dev2026010204.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-keras_nightly-3.14.0.dev2026010204.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
-keras_nightly-3.14.0.dev2026010204.dist-info/RECORD,,
+keras_nightly-3.14.0.dev2026010304.dist-info/METADATA,sha256=TRFo4Hl7iWbmC4VRcNIC6z3OZn-UvvN1VOu_XgjYt7I,6339
+keras_nightly-3.14.0.dev2026010304.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+keras_nightly-3.14.0.dev2026010304.dist-info/top_level.txt,sha256=ptcw_-QuGZ4ZDjMdwi_Z0clZm8QAqFdvzzFnDEOTs9o,6
+keras_nightly-3.14.0.dev2026010304.dist-info/RECORD,,