Add the bartlett function (#3155)

mlx/ops.cpp

@@ -2312,6 +2312,21 @@ array argmax(
   return out;
 }
 
+array bartlett(int M, StreamOrDevice s /* = {} */) {
+  if (M < 1) {
+    return array({});
+  }
+  if (M == 1) {
+    return ones({1}, float32, s);
+  }
+
+  auto n = arange(0, M, float32, s);
+  float factor_val = 2.0f / (M - 1);
+  auto factor = array(factor_val, float32);
+  auto term = subtract(multiply(factor, n, s), array(1.0f, float32), s);
+  return subtract(array(1.0f, float32), abs(term, s), s);
+}
+
 array hanning(int M, StreamOrDevice s /* = {} */) {
   if (M < 1) {
     return array({});

mlx/ops.h

@@ -672,6 +672,9 @@ MLX_API array hanning(int M, StreamOrDevice s = {});
 /** Returns the Hamming window of size M. */
 MLX_API array hamming(int M, StreamOrDevice s = {});
 
+/** Returns the bartlett window of size M. */
+MLX_API array bartlett(int M, StreamOrDevice s = {});
+
 /** Returns the Blackmann window of size M. */
 MLX_API array blackman(int M, StreamOrDevice s = {});
 

python/src/ops.cpp

@@ -1428,6 +1428,28 @@ void init_ops(nb::module_& m) {
       "stream"_a = nb::none(),
       nb::sig(
           "def arange(stop : Union[int, float], step : Union[None, int, float] = None, dtype: Optional[Dtype] = None, *, stream: Union[None, Stream, Device] = None) -> array"));
+  m.def(
+      "bartlett",
+      &mlx::core::bartlett,
+      "M"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"pbdoc(
+        Return the Bartlett window.
+        
+        The Bartlett window is a taper formed by using a weighted cosine.
+
+        .. math::
+          w(n) = 1 - \frac{2|n - (M-1)/2|}{M-1}
+           \qquad 0 \le n \le M-1
+        
+        Args:
+            M (int): Number of points in the output window.
+            
+        Returns:
+            array: The window, with the maximum value normalized to one (the value one
+                   appears only if the number of samples is odd).
+    )pbdoc");
   m.def(
       "hanning",
       &mlx::core::hanning,

python/tests/test_ops.py

@@ -1474,6 +1474,18 @@ class TestOps(mlx_tests.MLXTestCase):
         self.assertEqual(a.size, 0)
         self.assertEqual(a.dtype, mx.float32)
 
+    def test_bartlett_general(self):
+        a = mx.bartlett(10)
+        expected = np.bartlett(10)
+        self.assertTrue(np.allclose(a, expected, atol=1e-5))
+
+        a = mx.bartlett(1)
+        self.assertEqual(a.item(), 1.0)
+
+        a = mx.bartlett(0)
+        self.assertEqual(a.size, 0)
+        self.assertEqual(a.dtype, mx.float32)
+
     def test_blackman_general(self):
         a = mx.blackman(10)
         expected = np.blackman(10)