diff --git a/mlx/ops.cpp b/mlx/ops.cpp index b07ed595..f875562e 100644 --- a/mlx/ops.cpp +++ b/mlx/ops.cpp @@ -2324,6 +2324,27 @@ array hanning(int M, StreamOrDevice s /* = {} */) { return square(sin(multiply(factor, n, s), s), s); } +array hamming(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.0 * M_PI) / (M - 1); + auto factor = array(factor_val, float32); + + auto arg = multiply(factor, n, s); + auto cos_vals = cos(arg, s); + + auto left_coef = array(0.54f, float32); + auto right_coef = array(0.46f, float32); + + return subtract(left_coef, multiply(right_coef, cos_vals, s), s); +} + /** Returns a sorted copy of the flattened array. */ array sort(const array& a, StreamOrDevice s /* = {} */) { int size = a.size(); diff --git a/mlx/ops.h b/mlx/ops.h index d1068fe5..3be8352e 100644 --- a/mlx/ops.h +++ b/mlx/ops.h @@ -669,6 +669,9 @@ min(const array& a, int axis, bool keepdims = false, StreamOrDevice s = {}); /** Returns the Hanning window of size M. */ 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 index of the minimum value in the array. */ MLX_API array argmin(const array& a, bool keepdims, StreamOrDevice s = {}); inline array argmin(const array& a, StreamOrDevice s = {}) { diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 46fc8d84..f972d8b4 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -1451,6 +1451,30 @@ void init_ops(nb::module_& m) { appears only if the number of samples is odd). )pbdoc"); m.def( + "hamming", + &mlx::core::hamming, + "M"_a, + nb::kw_only(), + "stream"_a = nb::none(), + nb::sig( + "def hamming(M: int, *, stream: Union[None, Stream, Device] = None) -> array"), + R"pbdoc( + Return the Hamming window. + + The Hamming window is a taper formed by using a weighted cosine. + + .. math:: + w(n) = 0.54 - 0.46 \cos\left(\frac{2\pi n}{M-1}\right) + \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( "linspace", [](Scalar start, Scalar stop, diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py index 4365e873..bd8f5850 100644 --- a/python/tests/test_ops.py +++ b/python/tests/test_ops.py @@ -1462,6 +1462,18 @@ class TestOps(mlx_tests.MLXTestCase): self.assertEqual(a.size, 0) self.assertEqual(a.dtype, mx.float32) + def test_hamming_general(self): + a = mx.hamming(10) + expected = np.hamming(10) + self.assertTrue(np.allclose(a, expected, atol=1e-5)) + + a = mx.hamming(1) + self.assertEqual(a.item(), 1.0) + + a = mx.hamming(0) + self.assertEqual(a.size, 0) + self.assertEqual(a.dtype, mx.float32) + def test_unary_ops(self): def test_ops(npop, mlxop, x, y, atol, rtol): r_np = npop(x)