Add the bartlett function (#3155)

2026-03-03 11:40:54 -08:00
parent f145ece976
commit 8cd377b7db
4 changed files with 52 additions and 0 deletions
@@ -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)