Fix TypeError with empty axes in FFT functions

Author: antonwolfy

mkl_fft/_fft_utils.py

@@ -233,8 +233,10 @@ def _iter_complementary(x, axes, func, kwargs, result):
         m_ind = _flat_to_multi(ind, sub_shape)
         for k1, k2 in zip(dual_ind, m_ind):
             sl[k1] = k2
+        tsl = tuple(sl)
+
         if np.issubdtype(x.dtype, np.complexfloating):
-            func(x[tuple(sl)], **kwargs, out=result[tuple(sl)])
+            func(x[tsl], **kwargs, out=result[tsl])
         else:
             # For c2c FFT, if the input is real, half of the output is the
             # complex conjugate of the other half. Instead of upcasting the
@@ -247,7 +249,7 @@ def _iter_complementary(x, axes, func, kwargs, result):
             # array appeared in the second half of the NumPy output array,
             # while the equivalent element in the NumPy array was the conjugate
             # of the mkl_fft output array.
-            np.copyto(result[tuple(sl)], func(x[tuple(sl)], **kwargs))
+            np.copyto(result[tsl], func(x[tsl], **kwargs))
 
     return result
 
@@ -260,7 +262,6 @@ def _iter_fftnd(
     direction=+1,
     scale_function=lambda ind: 1.0,
 ):
-    a = np.asarray(a)
     s, axes = _init_nd_shape_and_axes(a, s, axes)
 
     # Combine the two, but in reverse, to end with the first axis given.
@@ -412,8 +413,20 @@ def _c2c_fftnd_impl(
             out=out,
         )
     else:
+        x = np.asarray(x)
+
+        # Fast path: FFT over no axes is identity (just type conversion, no scaling)
+        _, xa = _cook_nd_args(x, s, axes)
+        if len(xa) == 0:
+            if out is None:
+                out = x.astype(dtype=_output_dtype(x.dtype), copy=True)
+            else:
+                _validate_out_array(out, x, _output_dtype(x.dtype))
+                np.copyto(out, x)
+            # No scaling applied - identity transform has no normalization
+            return out
+
         if _complementary and x.dtype in valid_dtypes:
-            x = np.asarray(x)
             if out is None:
                 res = np.empty_like(x, dtype=_output_dtype(x.dtype))
             else:

mkl_fft/tests/test_fftnd.py

@@ -317,3 +317,48 @@ def test_out_strided(axes, func):
     expected = getattr(np.fft, func)(x, axes=axes, out=out)
 
     assert_allclose(result, expected, strict=True)
+
+
+@pytest.mark.parametrize(
+    "dtype", [np.float32, np.float64, np.complex64, np.complex128]
+)
+@pytest.mark.parametrize("shape", [(3, 4), (5,), (2, 3, 4), (10, 20)])
+@pytest.mark.parametrize("norm", [None, "ortho", "forward", "backward"])
+@pytest.mark.parametrize("func", ["fftn", "ifftn", "fft2", "ifft2"])
+def test_empty_axes(dtype, shape, norm, func):
+    if np.issubdtype(dtype, np.complexfloating):
+        x = rnd.random(shape).astype(dtype) + 1j * rnd.random(shape).astype(
+            dtype
+        )
+    else:
+        x = rnd.random(shape).astype(dtype)
+
+    # Test fftn with axes=()
+    result = getattr(mkl_fft, func)(x, axes=(), norm=norm)
+    expected = getattr(np.fft, func)(x, axes=(), norm=norm)
+
+    rtol, atol = _get_rtol_atol(result)
+    assert_allclose(result, expected, rtol=rtol, atol=atol, strict=True)
+
+
+@pytest.mark.parametrize(
+    "dtype", [np.float32, np.float64, np.complex64, np.complex128]
+)
+@pytest.mark.parametrize("func", ["fftn", "ifftn", "fft2", "ifft2"])
+def test_empty_axes_with_out(dtype, func):
+    if np.issubdtype(dtype, np.complexfloating):
+        x = rnd.random((3, 4)).astype(dtype) + 1j * rnd.random((3, 4)).astype(
+            dtype
+        )
+    else:
+        x = rnd.random((3, 4)).astype(dtype)
+
+    out_dtype = np.dtype(dtype).char.upper()
+    out = np.empty_like(x, dtype=out_dtype)
+    result = getattr(mkl_fft, func)(x, axes=(), out=out)
+    expected = getattr(np.fft, func)(x, axes=())
+
+    # Result should be written to out
+    assert result is out
+    rtol, atol = _get_rtol_atol(result)
+    assert_allclose(result, expected, rtol=rtol, atol=atol, strict=True)