zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-12-16 01:49:05 +08:00
Code Issues Actions 1 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

7
mlx/ops.cpp
View File

@@ -5025,8 +5025,11 @@ array roll(
}
auto sh = shift[i];
auto split_index =
(sh < 0) ? (-sh) % a.shape(ax) : a.shape(ax) - sh % a.shape(ax);
auto size = a.shape(ax);
if (size == 0) {
continue; // skip rolling this axis if it has size 0
}
auto split_index = (sh < 0) ? (-sh) % size : size - sh % size;
auto parts = split(result, Shape{split_index}, ax, s);
std::swap(parts[0], parts[1]);