keras-nightly 3.14.0.dev2025123104__tar.gz → 3.14.0.dev2026010404__tar.gz

{keras_nightly-3.14.0.dev2025123104 → keras_nightly-3.14.0.dev2026010404}/PKG-INFO

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

{keras_nightly-3.14.0.dev2025123104 → keras_nightly-3.14.0.dev2026010404}/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_nightly-3.14.0.dev2025123104 → keras_nightly-3.14.0.dev2026010404}/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_nightly-3.14.0.dev2025123104 → keras_nightly-3.14.0.dev2026010404}/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.dev2025123104"
+__version__ = "3.14.0.dev2026010404"
 
 
 @keras_export("keras.version")

{keras_nightly-3.14.0.dev2025123104 → keras_nightly-3.14.0.dev2026010404}/keras_nightly.egg-info/PKG-INFO

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