Implement bartlett function in keras.ops (#21214)

* Add bartlett for ops

* Update excluded_concrete_tests.txt

Author: shashaka

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

@@ -135,6 +135,7 @@
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
 from keras.src.ops.numpy import average as average
+from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount
 from keras.src.ops.numpy import bitwise_and as bitwise_and
 from keras.src.ops.numpy import bitwise_invert as bitwise_invert

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

@@ -27,6 +27,7 @@
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
 from keras.src.ops.numpy import average as average
+from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount
 from keras.src.ops.numpy import bitwise_and as bitwise_and
 from keras.src.ops.numpy import bitwise_invert as bitwise_invert

keras/api/ops/__init__.py

@@ -135,6 +135,7 @@
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
 from keras.src.ops.numpy import average as average
+from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount
 from keras.src.ops.numpy import bitwise_and as bitwise_and
 from keras.src.ops.numpy import bitwise_invert as bitwise_invert

keras/api/ops/numpy/__init__.py

@@ -27,6 +27,7 @@
 from keras.src.ops.numpy import argsort as argsort
 from keras.src.ops.numpy import array as array
 from keras.src.ops.numpy import average as average
+from keras.src.ops.numpy import bartlett as bartlett
 from keras.src.ops.numpy import bincount as bincount
 from keras.src.ops.numpy import bitwise_and as bitwise_and
 from keras.src.ops.numpy import bitwise_invert as bitwise_invert

keras/src/backend/jax/numpy.py

@@ -37,6 +37,11 @@ def add(x1, x2):
     return jnp.add(x1, x2)
 
 
+def bartlett(x):
+    x = convert_to_tensor(x)
+    return jnp.bartlett(x)
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     # Note: bincount is never tracable / jittable because the output shape
     # depends on the values in x.

keras/src/backend/numpy/numpy.py

@@ -305,6 +305,11 @@ def average(x, axis=None, weights=None):
     return np.average(x, weights=weights, axis=axis)
 
 
+def bartlett(x):
+    x = convert_to_tensor(x)
+    return np.bartlett(x).astype(config.floatx())
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with numpy backend")

keras/src/backend/openvino/excluded_concrete_tests.txt

@@ -8,6 +8,7 @@ NumpyDtypeTest::test_angle
 NumpyDtypeTest::test_any
 NumpyDtypeTest::test_argpartition
 NumpyDtypeTest::test_array
+NumpyDtypeTest::test_bartlett
 NumpyDtypeTest::test_bitwise
 NumpyDtypeTest::test_ceil
 NumpyDtypeTest::test_concatenate
@@ -75,6 +76,7 @@ NumpyOneInputOpsCorrectnessTest::test_angle
 NumpyOneInputOpsCorrectnessTest::test_any
 NumpyOneInputOpsCorrectnessTest::test_argpartition
 NumpyOneInputOpsCorrectnessTest::test_array
+NumpyOneInputOpsCorrectnessTest::test_bartlett
 NumpyOneInputOpsCorrectnessTest::test_bitwise_invert
 NumpyOneInputOpsCorrectnessTest::test_conj
 NumpyOneInputOpsCorrectnessTest::test_correlate
@@ -151,4 +153,5 @@ NumpyTwoInputOpsCorrectnessTest::test_tensordot
 NumpyTwoInputOpsCorrectnessTest::test_vdot
 NumpyTwoInputOpsCorrectnessTest::test_where
 NumpyOneInputOpsDynamicShapeTest::test_angle
+NumpyOneInputOpsDynamicShapeTest::test_bartlett
 NumpyOneInputOpsStaticShapeTest::test_angle

keras/src/backend/tensorflow/numpy.py

@@ -131,6 +131,21 @@ def add(x1, x2):
     return tf.add(x1, x2)
 
 
+def bartlett(x):
+    x = convert_to_tensor(x, dtype=config.floatx())
+    if x == 0:
+        return tf.constant([])
+    if x == 1:
+        return tf.ones([1])
+
+    n = tf.range(x)
+    half = (x - 1) / 2
+
+    window = tf.where(n <= half, 2.0 * n / (x - 1), 2.0 - 2.0 * n / (x - 1))
+
+    return window
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     x = convert_to_tensor(x)
     dtypes_to_resolve = [x.dtype]

keras/src/backend/torch/numpy.py

@@ -430,6 +430,11 @@ def average(x, axis=None, weights=None):
     return torch.mean(x, axis)
 
 
+def bartlett(x):
+    x = convert_to_tensor(x)
+    return torch.signal.windows.bartlett(x)
+
+
 def bincount(x, weights=None, minlength=0, sparse=False):
     if sparse:
         raise ValueError("Unsupported value `sparse=True` with torch backend")

keras/src/ops/numpy.py

@@ -1210,6 +1210,35 @@ def average(x, axis=None, weights=None):
     return backend.numpy.average(x, weights=weights, axis=axis)
 
 
+class Bartlett(Operation):
+    def call(self, x):
+        return backend.numpy.bartlett(x)
+
+    def compute_output_spec(self, x):
+        return KerasTensor(x.shape, dtype=backend.floatx())
+
+
+@keras_export(["keras.ops.bartlett", "keras.ops.numpy.bartlett"])
+def bartlett(x):
+    """Bartlett window function.
+    The Bartlett window is a triangular window that rises then falls linearly.
+
+    Args:
+        x: Scalar or 1D Tensor. Window length.
+
+    Returns:
+        A 1D tensor containing the Bartlett window values.
+
+    Example:
+    >>> x = keras.ops.convert_to_tensor(5)
+    >>> keras.ops.bartlett(x)
+    array([0. , 0.5, 1. , 0.5, 0. ], dtype=float32)
+    """
+    if any_symbolic_tensors((x,)):
+        return Bartlett().symbolic_call(x)
+    return backend.numpy.bartlett(x)
+
+
 class Bincount(Operation):
     def __init__(self, weights=None, minlength=0, sparse=False):
         super().__init__()