zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-11 15:06:42 +08:00
Code Issues Actions 2 Packages Projects Releases Wiki Activity

add fftshift and ifftshift fft helpers (#2135)

Browse Source 
* 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 <aashiq@aashiq-mbp-m4.local>
This commit is contained in:
Aashiq Dheeraj
2025-04-30 01:13:45 -04:00
committed by GitHub
parent 7bb063bcb3
commit bb6565ef14
9 changed files with 275 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
python/src/fft.cpp
View File

@@ -459,4 +459,55 @@ void init_fft(nb::module_& parent_module) {
Returns:
array: The real array containing the inverse of :func:`rfftn`.
)pbdoc");
m.def(
"fftshift",
[](const mx::array& a,
const std::optional<std::vector<int>>& axes,
mx::StreamOrDevice s) {
if (axes.has_value()) {
return mx::fft::fftshift(a, axes.value(), s);
} else {
return mx::fft::fftshift(a, s);
}
},
"a"_a,
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
Shift the zero-frequency component to the center of the spectrum.
Args:
a (array): The input array.
axes (list(int), optional): Axes over which to perform the shift.
If ``None``, shift all axes.
Returns:
array: The shifted array with the same shape as the input.
)pbdoc");
m.def(
"ifftshift",
[](const mx::array& a,
const std::optional<std::vector<int>>& axes,
mx::StreamOrDevice s) {
if (axes.has_value()) {
return mx::fft::ifftshift(a, axes.value(), s);
} else {
return mx::fft::ifftshift(a, s);
}
},
"a"_a,
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
The inverse of :func:`fftshift`. While identical to :func:`fftshift` for even-length axes,
the behavior differs for odd-length axes.
Args:
a (array): The input array.
axes (list(int), optional): Axes over which to perform the inverse shift.
If ``None``, shift all axes.
Returns:
array: The inverse-shifted array with the same shape as the input.
)pbdoc");
}

62
python/tests/test_fft.py
View File

@@ -199,6 +199,68 @@ class TestFFT(mlx_tests.MLXTestCase):
with self.assertRaises(ValueError):
mx.fft.irfftn(x)
def test_fftshift(self):
# Test 1D arrays
r = np.random.rand(100).astype(np.float32)
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r)
# Test with specific axis
r = np.random.rand(4, 6).astype(np.float32)
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[0])
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[1])
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[0, 1])
# Test with negative axes
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[-1])
# Test with odd lengths
r = np.random.rand(5, 7).astype(np.float32)
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r)
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, r, axes=[0])
# Test with complex input
r = np.random.rand(8, 8).astype(np.float32)
i = np.random.rand(8, 8).astype(np.float32)
c = r + 1j * i
self.check_mx_np(mx.fft.fftshift, np.fft.fftshift, c)
def test_ifftshift(self):
# Test 1D arrays
r = np.random.rand(100).astype(np.float32)
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r)
# Test with specific axis
r = np.random.rand(4, 6).astype(np.float32)
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[0])
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[1])
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[0, 1])
# Test with negative axes
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[-1])
# Test with odd lengths
r = np.random.rand(5, 7).astype(np.float32)
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r)
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, r, axes=[0])
# Test with complex input
r = np.random.rand(8, 8).astype(np.float32)
i = np.random.rand(8, 8).astype(np.float32)
c = r + 1j * i
self.check_mx_np(mx.fft.ifftshift, np.fft.ifftshift, c)
def test_fftshift_errors(self):
# Test invalid axes
x = mx.array(np.random.rand(4, 4).astype(np.float32))
with self.assertRaises(ValueError):
mx.fft.fftshift(x, axes=[2])
with self.assertRaises(ValueError):
mx.fft.fftshift(x, axes=[-3])
# Test empty array
x = mx.array([])
self.assertTrue(mx.array_equal(mx.fft.fftshift(x), x))
if __name__ == "__main__":
unittest.main()

5
python/tests/test_ops.py
View File

@@ -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)