add fftshift and ifftshift fft helpers (#2135)

* add fftshift and ifftshift fft helpers * address comments * axes have to be iterable * fix fp error in roll + add test --------- Co-authored-by: Aashiq Dheeraj <[email protected]>
2025-04-29 22:13:45 -07:00
parent 7bb063bcb3
commit bb6565ef14
9 changed files with 275 additions and 2 deletions
@@ -2961,6 +2961,11 @@ class TestOps(mlx_tests.MLXTestCase):
            y2 = mx.roll(x, s, a)
            self.assertTrue(mx.array_equal(y1, y2).item())

+    def test_roll_errors(self):
+        x = mx.array([])
+        result = mx.roll(x, [0], [0])
+        self.assertTrue(mx.array_equal(result, x))
+
    def test_real_imag(self):
        x = mx.random.uniform(shape=(4, 4))
        out = mx.real(x)