keras-nightly 3.14.0.dev2026012304__tar.gz → 3.14.0.dev2026012504__tar.gz

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/PKG-INFO

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

{keras_nightly-3.14.0.dev2026012304/keras/api → keras_nightly-3.14.0.dev2026012504/keras/api/_tf_keras/keras}/ops/__init__.py

@@ -245,6 +245,7 @@ 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 nanmax as nanmax
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012304/keras/api → keras_nightly-3.14.0.dev2026012504/keras/api/_tf_keras/keras}/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ 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 nanmax as nanmax
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012304/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012504/keras/api}/ops/__init__.py

@@ -245,6 +245,7 @@ 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 nanmax as nanmax
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012304/keras/api/_tf_keras/keras → keras_nightly-3.14.0.dev2026012504/keras/api}/ops/numpy/__init__.py

@@ -129,6 +129,7 @@ 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 nanmax as nanmax
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import ndim as ndim

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/jax/numpy.py

@@ -1013,6 +1013,11 @@ def moveaxis(x, source, destination):
     return jnp.moveaxis(x, source=source, destination=destination)
 
 
+def nanmax(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    return jnp.nanmax(x, axis=axis, keepdims=keepdims)
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     return jnp.nanmin(x, axis=axis, keepdims=keepdims)

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/numpy/numpy.py

@@ -960,6 +960,10 @@ def moveaxis(x, source, destination):
     return np.moveaxis(x, source=source, destination=destination)
 
 
+def nanmax(x, axis=None, keepdims=False):
+    return np.nanmax(x, axis=axis, keepdims=keepdims)
+
+
 def nanmin(x, axis=None, keepdims=False):
     return np.nanmin(x, axis=axis, keepdims=keepdims)
 

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/openvino/numpy.py

@@ -2107,6 +2107,10 @@ def moveaxis(x, source, destination):
     return OpenVINOKerasTensor(ov_opset.transpose(x, axes_const).output(0))
 
 
+def nanmax(x, axis=None, keepdims=False):
+    raise NotImplementedError("`nanmax` is not supported with openvino backend")
+
+
 def nanmin(x, axis=None, keepdims=False):
     raise NotImplementedError("`nanmin` is not supported with openvino backend")
 

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/tensorflow/numpy.py

@@ -2125,6 +2125,26 @@ def moveaxis(x, source, destination):
     return tf.transpose(x, perm)
 
 
+def nanmax(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+
+    if not x.dtype.is_floating:
+        dtype = standardize_dtype(x.dtype)
+        if dtype == "bool":
+            return tf.reduce_any(x, axis=axis, keepdims=keepdims)
+        return tf.reduce_max(x, axis=axis, keepdims=keepdims)
+
+    x_clean = tf.where(
+        tf.math.is_nan(x), tf.constant(float("-inf"), dtype=x.dtype), x
+    )
+
+    return tf.where(
+        tf.reduce_all(tf.math.is_nan(x), axis=axis, keepdims=keepdims),
+        tf.constant(float("nan"), dtype=x.dtype),
+        tf.reduce_max(x_clean, axis=axis, keepdims=keepdims),
+    )
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
 

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/torch/core.py

@@ -1,6 +1,7 @@
 import builtins
 import contextlib
 import functools
+import os
 
 import ml_dtypes
 import numpy as np
@@ -26,7 +27,9 @@ IS_THREAD_SAFE = True
 # Some operators such as 'aten::_foreach_mul_.Scalar'
 # are not currently implemented for the MPS device.
 # check https://github.com/pytorch/pytorch/issues/77764.
-if torch.backends.mps.is_available():
+if "KERAS_TORCH_DEVICE" in os.environ:
+    DEFAULT_DEVICE = os.environ["KERAS_TORCH_DEVICE"]
+elif torch.backends.mps.is_available():
     DEFAULT_DEVICE = "mps"
 elif torch.cuda.is_available():
     DEFAULT_DEVICE = "cuda"

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/torch/nn.py

@@ -1131,10 +1131,6 @@ def dot_product_attention(
     flash_attention=None,
     attn_logits_soft_cap=None,
 ):
-    if bias is not None:
-        raise ValueError(
-            "torch's `dot_product_attention` doesn't support `bias`."
-        )
     query = convert_to_tensor(query)
     key = convert_to_tensor(key)
     value = convert_to_tensor(value)
@@ -1144,6 +1140,10 @@ def dot_product_attention(
             f"Received: query.shape={query.shape}, key.shape={key.shape}, "
             f"value.shape={value.shape}."
         )
+    if bias is not None and mask is not None:
+        raise ValueError(
+            "Only one of `bias` and `mask` can be provided. Received both."
+        )
     compute_dtype = backend.result_type(query.dtype, key.dtype, value.dtype)
     query = cast(query, compute_dtype)
     key = cast(key, compute_dtype)
@@ -1154,6 +1154,9 @@ def dot_product_attention(
         # Explicit set `is_causal` to `False` when `mask` is not `None`.
         is_causal = False
         mask = torch.where(mask, 0.0, _get_large_negative(query.dtype))
+    if bias is not None:
+        bias = convert_to_tensor(bias, dtype=compute_dtype)
+        mask = bias  # Use `bias` as `mask` for scaled_dot_product_attention.
 
     axis0, axis1 = 1, 2
     query = torch.transpose(query, axis0, axis1)

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/backend/torch/numpy.py

@@ -1272,6 +1272,24 @@ def moveaxis(x, source, destination):
     return torch.moveaxis(x, source=source, destination=destination)
 
 
+def nanmax(x, axis=None, keepdims=False):
+    x = convert_to_tensor(x)
+    if not torch.is_floating_point(x):
+        return torch.amax(x, dim=axis, keepdim=keepdims)
+
+    if axis == () or axis == []:
+        return x
+
+    x_clean = torch.where(torch.isnan(x), float("-inf"), x)
+    out = torch.amax(x_clean, dim=axis, keepdim=keepdims)
+
+    return torch.where(
+        torch.isnan(x).all(dim=axis, keepdim=keepdims),
+        torch.tensor(float("nan"), dtype=x.dtype, device=get_device()),
+        out,
+    )
+
+
 def nanmin(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     if not torch.is_floating_point(x):

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/layers/rnn/time_distributed.py

@@ -71,37 +71,43 @@ class TimeDistributed(Wrapper):
         super().build(child_input_shape)
 
     def call(self, inputs, training=None, mask=None):
-        input_shape = ops.shape(inputs)
-        mask_shape = None if mask is None else ops.shape(mask)
-        batch_size = input_shape[0]
-        timesteps = input_shape[1]
-
-        # For TF backend with graph mode and `partial_batch_size`, skip
-        # evaluation of `batch_size` as it can be a `strided_slice` and
-        # not a constant.
-        if backend.backend() == "tensorflow":
-            from keras.src.utils.module_utils import tensorflow as tf
-
-            if (
-                not tf.executing_eagerly
-                and mask_shape is not None
-                and mask_shape[1:2] != (timesteps,)
-            ):
+        # Validate mask shape using static shape info when available
+        if mask is not None:
+            mask_shape = mask.shape
+            input_shape = inputs.shape
+
+            # Check if mask has at least 2 dimensions (batch and timesteps)
+            if len(mask_shape) < 2:
                 raise ValueError(
-                    "`TimeDistributed` Layer should be passed a `mask` of "
-                    f"shape ({batch_size}, {timesteps}, ...), "
-                    f"received: mask.shape={mask_shape}"
+                    "The `mask` passed to the `TimeDistributed` layer must be "
+                    "at least 2D (e.g., `(batch_size, timesteps)`), but it has "
+                    f"{len(mask_shape)} dimension(s) with shape {mask_shape}."
                 )
-        elif mask_shape is not None and mask_shape[:2] != (
-            batch_size,
-            timesteps,
-        ):
-            raise ValueError(
-                "`TimeDistributed` Layer should be passed a `mask` of "
-                f"shape ({batch_size}, {timesteps}, ...), "
-                f"received: mask.shape={mask_shape}"
+
+            # Check batch size and timesteps dimensions match
+            batch_mismatch = (
+                input_shape[0] is not None
+                and mask_shape[0] is not None
+                and input_shape[0] != mask_shape[0]
+            )
+            time_mismatch = (
+                input_shape[1] is not None
+                and mask_shape[1] is not None
+                and input_shape[1] != mask_shape[1]
             )
 
+            if batch_mismatch or time_mismatch:
+                raise ValueError(
+                    "The `mask` passed to the `TimeDistributed` layer has a "
+                    f"shape {mask_shape} that is incompatible with the input "
+                    f"shape {input_shape}. The first two dimensions of the "
+                    "mask (batch size and timesteps) must match the input's "
+                    "first two dimensions. Expected mask shape prefix: "
+                    f"({input_shape[0]}, {input_shape[1]})."
+                )
+
+        input_shape = ops.shape(inputs)
+
         def time_distributed_transpose(data):
             """Swaps the timestep and batch dimensions of a tensor."""
             axes = [1, 0, *range(2, len(data.shape))]
@@ -129,5 +135,7 @@ class TimeDistributed(Wrapper):
 
         # Implementation #2: use backend.vectorized_map.
 
-        outputs = backend.vectorized_map(step_function, ops.arange(timesteps))
+        outputs = backend.vectorized_map(
+            step_function, ops.arange(input_shape[0])
+        )
         return time_distributed_transpose(outputs)

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/ops/core.py

@@ -277,7 +277,7 @@ def associative_scan(f, elems, reverse=False, axis=0):
     [0, 1, 3, 6, 10]
 
     >>> sum_fn = lambda x, y: [x[0] + y[0], x[1] + y[1], x[2] + y[2]]
-    >>> xs = [keras.ops.array([[1, 2]]) for _ in range(3)]
+    >>> xs = [keras.ops.array([1, 2]) for _ in range(3)]
     >>> ys = keras.ops.associative_scan(sum_fn, xs, axis=0)
     >>> ys
     [[1, 3], [1, 3], [1, 3]]

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras/src/ops/numpy.py

@@ -5064,6 +5064,63 @@ def moveaxis(x, source, destination):
     return backend.numpy.moveaxis(x, source=source, destination=destination)
 
 
+class Nanmax(Operation):
+    def __init__(self, axis=None, keepdims=False, *, name=None):
+        super().__init__(name=name)
+        self.axis = axis
+        self.keepdims = keepdims
+
+    def call(self, x):
+        return backend.numpy.nanmax(x, axis=self.axis, keepdims=self.keepdims)
+
+    def compute_output_spec(self, x):
+        dtype = dtypes.result_type(getattr(x, "dtype", backend.floatx()))
+
+        if backend.backend() == "torch" and dtype == "uint32":
+            dtype = "int32"
+
+        return KerasTensor(
+            reduce_shape(x.shape, axis=self.axis, keepdims=self.keepdims),
+            dtype=dtype,
+        )
+
+
+@keras_export(["keras.ops.nanmax", "keras.ops.numpy.nanmax"])
+def nanmax(x, axis=None, keepdims=False):
+    """Maximum of a tensor over the given axes, ignoring NaNs.
+
+    Args:
+        x: Input tensor.
+        axis: Axis or axes along which the maximum is computed.
+            The default is to compute the maximum of the flattened tensor.
+        keepdims: If this is set to `True`, the axes which are reduced are left
+            in the result as dimensions with size one. Defaults
+            to `False`.
+
+    Returns:
+        Output tensor containing the maximum, with NaN values ignored. If all
+        values along a reduced axis are NaN, the result is NaN.
+
+    Examples:
+    >>> import numpy as np
+    >>> from keras import ops
+    >>> x = np.array([[1.0, np.nan, 3.0],
+    ...               [np.nan, 2.0, 1.0]])
+    >>> ops.nanmax(x)
+    3.0
+
+    >>> ops.nanmax(x, axis=1)
+    array([3., 2.])
+
+    >>> ops.nanmax(x, axis=1, keepdims=True)
+    array([[3.],
+           [2.]])
+    """
+    if any_symbolic_tensors((x,)):
+        return Nanmax(axis=axis, keepdims=keepdims).symbolic_call(x)
+    return backend.numpy.nanmax(x, axis=axis, keepdims=keepdims)
+
+
 class Nanmin(Operation):
     def __init__(self, axis=None, keepdims=False, *, name=None):
         super().__init__(name=name)

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

{keras_nightly-3.14.0.dev2026012304 → keras_nightly-3.14.0.dev2026012504}/keras_nightly.egg-info/PKG-INFO

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