From 3bbe87e6dcc7797add541a5b2ff39b7e76c8e90b Mon Sep 17 00:00:00 2001
From: willem adnet <140308890+Vlor999@users.noreply.github.com>
Date: Mon, 16 Feb 2026 18:44:49 +0100
Subject: [PATCH] Add hanning window function (#3124)

---
 mlx/ops.cpp              | 13 +++++++++++++
 mlx/ops.h                |  3 +++
 python/src/ops.cpp       | 22 ++++++++++++++++++++++
 python/tests/test_ops.py | 12 ++++++++++++
 4 files changed, 50 insertions(+)

diff --git a/mlx/ops.cpp b/mlx/ops.cpp
index 076de699..b07ed595 100644
--- a/mlx/ops.cpp
+++ b/mlx/ops.cpp
@@ -2311,6 +2311,19 @@ array argmax(
   return out;
 }
 
+array hanning(int M, StreamOrDevice s /* = {} */) {
+  if (M < 1) {
+    return array({});
+  }
+  if (M == 1) {
+    return ones({1}, float32, s);
+  }
+
+  auto n = arange(0, M, float32, s);
+  array factor(M_PI / (M - 1), float32);
+  return square(sin(multiply(factor, n, s), 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 1ff3bbfa..d1068fe5 100644
--- a/mlx/ops.h
+++ b/mlx/ops.h
@@ -666,6 +666,9 @@ min(const array& a,
 MLX_API array
 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 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 de51ad50..46fc8d84 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(
+      "hanning",
+      &mlx::core::hanning,
+      "M"_a,
+      nb::kw_only(),
+      "stream"_a = nb::none(),
+      R"pbdoc(
+        Return the Hanning window.
+        
+        The Hanning window is a taper formed by using a weighted cosine.
+
+        .. math::
+          w(n) = 0.5 - 0.5 \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,
diff --git a/python/tests/test_ops.py b/python/tests/test_ops.py
index ecdd970c..4365e873 100644
--- a/python/tests/test_ops.py
+++ b/python/tests/test_ops.py
@@ -1450,6 +1450,18 @@ class TestOps(mlx_tests.MLXTestCase):
         expected = [0]
         self.assertListEqual(a.tolist(), expected)
 
+    def test_hanning_general(self):
+        a = mx.hanning(10)
+        expected = np.hanning(10)
+        self.assertTrue(np.allclose(a, expected, atol=1e-5))
+
+        a = mx.hanning(1)
+        self.assertEqual(a.item(), 1.0)
+
+        a = mx.hanning(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)