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 <aashiq@aashiq-mbp-m4.local>

mlx/fft.cpp

@@ -184,8 +184,79 @@ array irfftn(
     StreamOrDevice s /* = {} */) {
   return fft_impl(a, axes, true, true, s);
 }
+
 array irfftn(const array& a, StreamOrDevice s /* = {} */) {
   return fft_impl(a, true, true, s);
 }
 
+array fftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[fftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    shifts.push_back(a.shape(axis) / 2);
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+array ifftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s /* = {} */) {
+  if (axes.empty()) {
+    return a;
+  }
+
+  Shape shifts;
+  for (int ax : axes) {
+    // Convert negative axes to positive
+    int axis = ax < 0 ? ax + a.ndim() : ax;
+    if (axis < 0 || axis >= a.ndim()) {
+      std::ostringstream msg;
+      msg << "[ifftshift] Invalid axis " << ax << " for array with " << a.ndim()
+          << " dimensions.";
+      throw std::invalid_argument(msg.str());
+    }
+    // Match NumPy's implementation
+    int size = a.shape(axis);
+    shifts.push_back(-(size / 2));
+  }
+
+  return roll(a, shifts, axes, s);
+}
+
+// Default versions that operate on all axes
+array fftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return fftshift(a, axes, s);
+}
+
+array ifftshift(const array& a, StreamOrDevice s /* = {} */) {
+  if (a.ndim() < 1) {
+    return a;
+  }
+  std::vector<int> axes(a.ndim());
+  std::iota(axes.begin(), axes.end(), 0);
+  return ifftshift(a, axes, s);
+}
+
 } // namespace mlx::core::fft

mlx/fft.h

@@ -145,5 +145,23 @@ inline array irfft2(
     StreamOrDevice s = {}) {
   return irfftn(a, axes, s);
 }
+/** Shift the zero-frequency component to the center of the spectrum. */
+array fftshift(const array& a, StreamOrDevice s = {});
+
+/** Shift the zero-frequency component to the center of the spectrum along
+ * specified axes. */
+array fftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
+
+/** The inverse of fftshift. */
+array ifftshift(const array& a, StreamOrDevice s = {});
+
+/** The inverse of fftshift along specified axes. */
+array ifftshift(
+    const array& a,
+    const std::vector<int>& axes,
+    StreamOrDevice s = {});
 
 } // namespace mlx::core::fft

mlx/ops.cpp

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