diff --git a/mlx/ops.cpp b/mlx/ops.cpp index 0cafd479..c7af8834 100644 --- a/mlx/ops.cpp +++ b/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({}); diff --git a/mlx/ops.h b/mlx/ops.h index 1f8331c1..74032c01 100644 --- a/mlx/ops.h +++ b/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 = {}); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 19231700..a4ce55f8 100644 --- a/python/src/ops.cpp +++ b/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, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index a4981161..fd40f3b6 100644 --- a/python/tests/test_ops.py +++ b/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)