zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-10-19 00:04:41 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

fix fft bug (#2062)

Browse Source
This commit is contained in:
Awni Hannun
2025-04-10 19:41:27 -07:00
committed by GitHub
parent ddaa4b7dcb
commit ef7ece9851
3 changed files with 52 additions and 29 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
python/tests/test_fft.py
View File

@@ -182,6 +182,18 @@ class TestFFT(mlx_tests.MLXTestCase):
out_np = np.abs(np.fft.fft(np.tile(np.reshape(np.array(b_np), (1, 4)), (4, 1))))
np.testing.assert_allclose(out_mx, out_np, atol=1e-5, rtol=1e-5)
def test_fft_into_ifft(self):
n_fft = 8193
mx.random.seed(0)
segment = mx.random.normal(shape=[1, n_fft]) + 1j * mx.random.normal(
shape=(1, n_fft)
)
segment = mx.fft.fft(segment, n=n_fft)
r = mx.fft.ifft(segment, n=n_fft)
r_np = np.fft.ifft(segment, n=n_fft)
self.assertTrue(np.allclose(r, r_np, atol=1e-5, rtol=1e-5))
if __name__ == "__main__":
unittest.main()