zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 09:21:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

address comments

Browse Source
This commit is contained in:
Aashiq Dheeraj 2025-04-29 20:42:27 -04:00
parent 00e43d18ed
commit 7c9746a0bc
2 changed files with 6 additions and 8 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
mlx/fft.cpp
View File

@ -207,7 +207,7 @@ array fftshift(
<< " dimensions.";
throw std::invalid_argument(msg.str());
}
// Match PyTorch's implementation
// Match NumPy's implementation
shifts.push_back(a.shape(axis) / 2);
}
@ -232,9 +232,9 @@ array ifftshift(
<< " dimensions.";
throw std::invalid_argument(msg.str());
}
// Match PyTorch's implementation
// Match NumPy's implementation
int size = a.shape(axis);
shifts.push_back((size + 1) / 2);
shifts.push_back(-(size / 2));
}
return roll(a, shifts, axes, s);

8
python/src/fft.cpp
View File

@ -479,8 +479,7 @@ void init_fft(nb::module_& parent_module) {
Args:
a (array): The input array.
axes (list(int), optional): Axes over which to perform the shift.
If None, shift all axes. Each axis can be negative and will be
converted to a positive axis using the same rules as NumPy.
If ``None``, shift all axes.
Returns:
array: The shifted array with the same shape as the input.
@ -500,14 +499,13 @@ void init_fft(nb::module_& parent_module) {
"axes"_a = nb::none(),
"stream"_a = nb::none(),
R"pbdoc(
The inverse of fftshift. While identical to fftshift for even-length axes,
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. Each axis can be negative and will be
converted to a positive axis using the same rules as NumPy.
If ``None``, shift all axes.
Returns:
array: The inverse-shifted array with the same shape as the input.