Implement nanquantile function in keras.ops

keras/api/_tf_keras/keras/ops/__init__.py

@@ -260,6 +260,7 @@
 from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nanprod as nanprod
+from keras.src.ops.numpy import nanquantile as nanquantile
 from keras.src.ops.numpy import nanstd as nanstd
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import nanvar as nanvar

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

@@ -141,6 +141,7 @@
 from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nanprod as nanprod
+from keras.src.ops.numpy import nanquantile as nanquantile
 from keras.src.ops.numpy import nanstd as nanstd
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import nanvar as nanvar

keras/api/ops/__init__.py

@@ -260,6 +260,7 @@
 from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nanprod as nanprod
+from keras.src.ops.numpy import nanquantile as nanquantile
 from keras.src.ops.numpy import nanstd as nanstd
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import nanvar as nanvar

keras/api/ops/numpy/__init__.py

@@ -141,6 +141,7 @@
 from keras.src.ops.numpy import nanmean as nanmean
 from keras.src.ops.numpy import nanmin as nanmin
 from keras.src.ops.numpy import nanprod as nanprod
+from keras.src.ops.numpy import nanquantile as nanquantile
 from keras.src.ops.numpy import nanstd as nanstd
 from keras.src.ops.numpy import nansum as nansum
 from keras.src.ops.numpy import nanvar as nanvar

keras/src/backend/jax/numpy.py

@@ -1083,6 +1083,12 @@ def nanprod(x, axis=None, keepdims=False):
     return jnp.nanprod(x, axis=axis, keepdims=keepdims)
 
 
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    x = convert_to_tensor(x)
+    q = convert_to_tensor(q)
+    return jnp.nanquantile(x, q, axis=axis, method=method, keepdims=keepdims)
+
+
 def nanstd(x, axis=None, keepdims=False):
     x = convert_to_tensor(x)
     return jnp.nanstd(x, axis=axis, keepdims=keepdims)

keras/src/backend/numpy/numpy.py

@@ -1058,6 +1058,17 @@ def nanprod(x, axis=None, keepdims=False):
     return np.nanprod(x, axis=axis, keepdims=keepdims, dtype=dtype)
 
 
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    x = convert_to_tensor(x)
+    ori_dtype = standardize_dtype(x.dtype)
+    if ori_dtype == "bool":
+        x = x.astype(config.floatx())
+    dtype = dtypes.result_type(x.dtype, float)
+    return np.nanquantile(
+        x, q, axis=axis, method=method, keepdims=keepdims
+    ).astype(dtype)
+
+
 def nanstd(x, axis=None, keepdims=False):
     axis = standardize_axis_for_numpy(axis)
     x = convert_to_tensor(x)

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -321,6 +321,7 @@ NumpyDtypeTest::test_maximum_python_types
 NumpyDtypeTest::test_minimum_python_types
 NumpyDtypeTest::test_nanargmax
 NumpyDtypeTest::test_nanargmin
+NumpyDtypeTest::test_nanquantile
 NumpyDtypeTest::test_power
 NumpyDtypeTest::test_sinc
 NumpyDtypeTest::test_view
@@ -339,6 +340,7 @@ NumpyOneInputOpsDynamicShapeTest::test_sinc
 NumpyOneInputOpsDynamicShapeTest::test_view
 NumpyOneInputOpsStaticShapeTest::test_sinc
 NumpyOneInputOpsStaticShapeTest::test_view
+NumpyTwoInputOpsCorrectnessTest::test_nanquantile
 OptimizerTest::test_constraints_are_applied
 OptimizerTest::test_ema
 OptimizerTest::test_gradient_accumulation

keras/src/backend/openvino/numpy.py

@@ -3012,6 +3012,12 @@ def nanprod(x, axis=None, keepdims=False):
     return OpenVINOKerasTensor(result)
 
 
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    raise NotImplementedError(
+        "`nanquantile` is not supported with openvino backend"
+    )
+
+
 def nanstd(x, axis=None, keepdims=False):
     return sqrt(nanvar(x, axis=axis, keepdims=keepdims))
 

keras/src/backend/tensorflow/numpy.py

@@ -2308,6 +2308,72 @@ def nanprod(x, axis=None, keepdims=False):
     return prod(x_safe, axis=axis, keepdims=keepdims)
 
 
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    x = convert_to_tensor(x)
+    q = convert_to_tensor(q, dtype=config.floatx())
+
+    def _nanquantile_1d(v):
+        valid = tf.boolean_mask(v, ~tf.math.is_nan(cast(v, config.floatx())))
+        return tf.cond(
+            tf.size(valid) > 0,
+            lambda: quantile(valid, q, method=method, keepdims=False),
+            lambda: tf.constant(float("nan"), dtype=x.dtype),
+        )
+
+    if axis is None:
+        x_flat = tf.reshape(x, [-1])
+        result = _nanquantile_1d(x_flat)
+
+        if keepdims:
+            new_shape = tf.concat(
+                [tf.shape(result), tf.ones(tf.rank(x), dtype=tf.int32)], axis=0
+            )
+            result = tf.reshape(result, new_shape)
+
+        return result
+
+    if isinstance(axis, int):
+        axis = [axis]
+    elif isinstance(axis, tuple):
+        axis = list(axis)
+
+    ndims = x.shape.rank
+    axis = [a if a >= 0 else a + ndims for a in axis]
+    other_axes = [i for i in range(ndims) if i not in axis]
+
+    perm = other_axes + axis
+    x_t = tf.transpose(x, perm)
+
+    shape = tf.shape(x_t)
+    other_rank = len(other_axes)
+    other_shape = shape[:other_rank]
+    reduce_shape = shape[other_rank:]
+
+    batch_size = tf.reduce_prod(other_shape)
+    reduction_size = tf.reduce_prod(reduce_shape)
+    x_flat = tf.reshape(x_t, [batch_size, reduction_size])
+
+    q_shape = tf.shape(q)
+
+    results = tf.map_fn(
+        _nanquantile_1d,
+        x_flat,
+        fn_output_signature=tf.TensorSpec(shape=q_shape, dtype=x.dtype),
+    )
+
+    if tf.rank(q) > 0:
+        results = tf.transpose(results)
+        results = tf.reshape(results, tf.concat([q_shape, other_shape], axis=0))
+    else:
+        results = tf.reshape(results, other_shape)
+
+    if keepdims:
+        for ax in sorted(axis):
+            results = tf.expand_dims(results, axis=ax + tf.rank(q))
+
+    return results
+
+
 def nanstd(x, axis=None, keepdims=False):
     var_val = nanvar(x, axis=axis, keepdims=keepdims)
     return tf.sqrt(var_val)

keras/src/backend/torch/numpy.py

@@ -1430,6 +1430,45 @@ def nanprod(x, axis=None, keepdims=False):
     )
 
 
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    x = convert_to_tensor(x)
+    q = convert_to_tensor(q)
+    axis = to_tuple_or_list(axis)
+
+    compute_dtype = dtypes.result_type(x.dtype, "float32")
+    result_dtype = dtypes.result_type(x.dtype, float)
+
+    x = cast(x, compute_dtype)
+    if x.dtype != q.dtype:
+        q = cast(q, x.dtype)
+
+    if axis is None:
+        y = reshape(x, [-1])
+    else:
+        axis = [canonicalize_axis(a, x.ndim) for a in axis]
+        other_dims = sorted(set(range(x.ndim)).difference(axis))
+        x_permed = torch.permute(x, dims=(other_dims + list(axis)))
+
+        x_shape = list(x.shape)
+        other_shape = [x_shape[i] for i in other_dims]
+        end_shape = [math.prod([x_shape[i] for i in axis])]
+        full_shape = other_shape + end_shape
+        y = reshape(x_permed, full_shape)
+
+    y = torch.nanquantile(y, q, dim=-1, interpolation=method)
+
+    if keepdims:
+        if axis is None:
+            for _ in range(x.ndim):
+                y = expand_dims(y, axis=-1)
+        else:
+            for i in sorted(axis):
+                i = i + 1 if q.ndim > 0 else i
+                y = expand_dims(y, axis=i)
+
+    return cast(y, result_dtype)
+
+
 def nanstd(x, axis=None, keepdims=False):
     var_val = nanvar(x, axis=axis, keepdims=keepdims)
     return torch.sqrt(var_val)

keras/src/ops/numpy.py

@@ -5792,6 +5792,90 @@ def nanprod(x, axis=None, keepdims=False):
     return backend.numpy.nanprod(x, axis=axis, keepdims=keepdims)
 
 
+class Nanquantile(Operation):
+    def __init__(
+        self, axis=None, method="linear", keepdims=False, *, name=None
+    ):
+        super().__init__(name=name)
+        self.axis = axis
+        self.method = method
+        self.keepdims = keepdims
+
+    def call(self, x, q):
+        return backend.numpy.nanquantile(
+            x, q, axis=self.axis, method=self.method, keepdims=self.keepdims
+        )
+
+    def compute_output_spec(self, x, q):
+        output_shape = reduce_shape(
+            x.shape, axis=self.axis, keepdims=self.keepdims
+        )
+        if hasattr(q, "shape"):
+            if len(q.shape) > 0:
+                output_shape = (q.shape[0],) + output_shape
+
+        if backend.standardize_dtype(x.dtype) == "int64":
+            dtype = backend.floatx()
+        else:
+            dtype = dtypes.result_type(x.dtype, float)
+
+        return KerasTensor(output_shape, dtype=dtype)
+
+
+@keras_export(["keras.ops.nanquantile", "keras.ops.numpy.nanquantile"])
+def nanquantile(x, q, axis=None, method="linear", keepdims=False):
+    """Compute the q-th quantile(s) of the data along the specified axis,
+    while ignoring NaNs.
+
+    Args:
+        x: Input tensor.
+        q: Probability or sequence of probabilities for the quantiles to
+            compute. Values must be between 0 and 1 inclusive.
+        axis: Axis or axes along which the quantiles are computed. Defaults to
+            `axis=None` which is to compute the quantile(s) along a flattened
+            version of the array.
+        method: A string specifies the method to use for estimating the
+            quantile. Available methods are `"linear"`, `"lower"`, `"higher"`,
+            `"midpoint"`, and `"nearest"`. Defaults to `"linear"`.
+            If the desired quantile lies between two data points `i < j`:
+            - `"linear"`: `i + (j - i) * fraction`, where fraction is the
+                fractional part of the index surrounded by `i` and `j`.
+            - `"lower"`: `i`.
+            - `"higher"`: `j`.
+            - `"midpoint"`: `(i + j) / 2`
+            - `"nearest"`: `i` or `j`, whichever is nearest.
+        keepdims: If this is set to `True`, the axes which are reduced
+            are left in the result as dimensions with size one.
+
+    Returns:
+        The quantile(s) ignoring NaNs.
+    Examples:
+    >>> import keras
+    >>> from keras import ops
+    >>> x = keras.ops.array([1., 2., 3., 4.])
+    >>> keras.ops.nanquantile(x, 0.5)
+    2.5
+    >>> x = keras.ops.array([1., 2., float("nan"), 4.])
+    >>> keras.ops.nanquantile(x, 0.5)
+    2.0
+    >>> x = keras.ops.array([1., 2., 3., 4.])
+    >>> keras.ops.nanquantile(x, [0.25, 0.75])
+    array([1.75, 3.25])
+    >>> x = keras.ops.array([[1., 2., float("nan")],
+    ...                      [4., 5., 6.]])
+    >>> keras.ops.nanquantile(x, 0.5, axis=1)
+    array([1.5, 5.0])
+    """
+    if any_symbolic_tensors((x, q)):
+        return Nanquantile(
+            axis=axis, method=method, keepdims=keepdims
+        ).symbolic_call(x, q)
+
+    return backend.numpy.nanquantile(
+        x, q, axis=axis, method=method, keepdims=keepdims
+    )
+
+
 class Nanstd(Operation):
     def __init__(self, axis=None, keepdims=False, *, name=None):
         super().__init__(name=name)