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

keras/_tf_keras/keras/dtype_policies/__init__.py

@@ -7,6 +7,9 @@ since your modifications would be overwritten.
 from keras.src.dtype_policies import deserialize as deserialize
 from keras.src.dtype_policies import get as get
 from keras.src.dtype_policies import serialize as serialize
+from keras.src.dtype_policies.dtype_policy import (
+    AWQDTypePolicy as AWQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import DTypePolicy as DTypePolicy
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,

keras/_tf_keras/keras/ops/__init__.py

@@ -245,8 +245,10 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
+from keras.src.ops.numpy import nextafter as nextafter
 from keras.src.ops.numpy import nonzero as nonzero
 from keras.src.ops.numpy import not_equal as not_equal
 from keras.src.ops.numpy import ones as ones
@@ -255,6 +257,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/_tf_keras/keras/ops/numpy/__init__.py

@@ -129,8 +129,10 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
+from keras.src.ops.numpy import nextafter as nextafter
 from keras.src.ops.numpy import nonzero as nonzero
 from keras.src.ops.numpy import not_equal as not_equal
 from keras.src.ops.numpy import ones as ones
@@ -139,6 +141,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/_tf_keras/keras/quantizers/__init__.py

@@ -7,6 +7,7 @@ since your modifications would be overwritten.
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
+from keras.src.quantizers.awq_config import AWQConfig as AWQConfig
 from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantization_config import (
     Float8QuantizationConfig as Float8QuantizationConfig,

keras/dtype_policies/__init__.py

@@ -7,6 +7,9 @@ since your modifications would be overwritten.
 from keras.src.dtype_policies import deserialize as deserialize
 from keras.src.dtype_policies import get as get
 from keras.src.dtype_policies import serialize as serialize
+from keras.src.dtype_policies.dtype_policy import (
+    AWQDTypePolicy as AWQDTypePolicy,
+)
 from keras.src.dtype_policies.dtype_policy import DTypePolicy as DTypePolicy
 from keras.src.dtype_policies.dtype_policy import (
     FloatDTypePolicy as FloatDTypePolicy,

keras/ops/__init__.py

@@ -245,8 +245,10 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
+from keras.src.ops.numpy import nextafter as nextafter
 from keras.src.ops.numpy import nonzero as nonzero
 from keras.src.ops.numpy import not_equal as not_equal
 from keras.src.ops.numpy import ones as ones
@@ -255,6 +257,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/ops/numpy/__init__.py

@@ -129,8 +129,10 @@ from keras.src.ops.numpy import mod as mod
 from keras.src.ops.numpy import moveaxis as moveaxis
 from keras.src.ops.numpy import multiply as multiply
 from keras.src.ops.numpy import nan_to_num as nan_to_num
+from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim
 from keras.src.ops.numpy import negative as negative
+from keras.src.ops.numpy import nextafter as nextafter
 from keras.src.ops.numpy import nonzero as nonzero
 from keras.src.ops.numpy import not_equal as not_equal
 from keras.src.ops.numpy import ones as ones
@@ -139,6 +141,7 @@ from keras.src.ops.numpy import outer as outer
 from keras.src.ops.numpy import pad as pad
 from keras.src.ops.numpy import power as power
 from keras.src.ops.numpy import prod as prod
+from keras.src.ops.numpy import ptp as ptp
 from keras.src.ops.numpy import quantile as quantile
 from keras.src.ops.numpy import ravel as ravel
 from keras.src.ops.numpy import real as real

keras/quantizers/__init__.py

@@ -7,6 +7,7 @@ since your modifications would be overwritten.
 from keras.src.quantizers import deserialize as deserialize
 from keras.src.quantizers import get as get
 from keras.src.quantizers import serialize as serialize
+from keras.src.quantizers.awq_config import AWQConfig as AWQConfig
 from keras.src.quantizers.gptq_config import GPTQConfig as GPTQConfig
 from keras.src.quantizers.quantization_config import (
     Float8QuantizationConfig as Float8QuantizationConfig,

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/backend/jax/numpy.py

@@ -1013,6 +1013,11 @@ def moveaxis(x, source, destination):
     return jnp.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.nansum(x, axis=axis, keepdims=keepdims)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     x = convert_to_tensor(x)
     return jnp.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
@@ -1063,6 +1068,11 @@ def prod(x, axis=None, keepdims=False, dtype=None):
     return jnp.prod(x, axis=axis, keepdims=keepdims, dtype=dtype)
 
 
+def ptp(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.ptp(x, axis=axis, keepdims=keepdims)
+
+
 def quantile(x, q, axis=None, method="linear", keepdims=False):
     x = convert_to_tensor(x)
     q = convert_to_tensor(q)
@@ -1351,6 +1361,12 @@ def negative(x):
     return jnp.negative(x)
 
 
+def nextafter(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    return jnp.nextafter(x1, x2)
+
+
 @sparse.elementwise_unary(linear=False)
 def square(x):
     x = convert_to_tensor(x)

keras/src/backend/numpy/numpy.py

@@ -960,6 +960,17 @@ def moveaxis(x, source, destination):
     return np.moveaxis(x, source=source, destination=destination)
 
 
+def nansum(x, axis=None, keepdims=False):
+    axis = standardize_axis_for_numpy(axis)
+    dtype = standardize_dtype(x.dtype)
+
+    if dtype in ("bool", "int8", "int16"):
+        dtype = "int32"
+    elif dtype in ("uint8", "uint16"):
+        dtype = "uint32"
+    return np.nansum(x, axis=axis, keepdims=keepdims).astype(dtype)
+
+
 def nan_to_num(x, nan=0.0, posinf=None, neginf=None):
     return np.nan_to_num(x, nan=nan, posinf=posinf, neginf=neginf)
 
@@ -1018,6 +1029,10 @@ def prod(x, axis=None, keepdims=False, dtype=None):
     return np.prod(x, axis=axis, keepdims=keepdims, dtype=dtype)
 
 
+def ptp(x, axis=None, keepdims=False):
+    return np.ptp(x, axis=axis, keepdims=keepdims)
+
+
 def quantile(x, q, axis=None, method="linear", keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)
@@ -1335,6 +1350,14 @@ def negative(x):
     return np.negative(x)
 
 
+def nextafter(x1, x2):
+    x1 = convert_to_tensor(x1)
+    x2 = convert_to_tensor(x2)
+    dtype = dtypes.result_type(x1.dtype, x2.dtype, float)
+
+    return np.nextafter(x1, x2).astype(dtype)
+
+
 def square(x):
     x = convert_to_tensor(x)
     if standardize_dtype(x.dtype) == "bool":