Add the hamming window function (#3135)

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();

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 = {}) {

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,

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)