reverse the order of individual FFTs in rfftn (#2524)

For `numpy.fft.fftn` and `numpy.fft.ifftn`, individual FFTs over `axes`
are performed in [reverse
order](https://github.com/numpy/numpy/blob/v2.2.0/numpy/fft/_pocketfft.py#L739).

Similarly, for `numpy.fft.rfftn`, individual FFTs are performed in
[reverse
order](https://github.com/numpy/numpy/blob/v2.2.0/numpy/fft/_pocketfft.py#L1382).

However, for `numpy.fft.irfftn` individual FFTs are performed in
[forward
order](https://github.com/numpy/numpy/blob/v2.2.0/numpy/fft/_pocketfft.py#L1600).

Author: vtavana

CHANGELOG.md

@@ -24,6 +24,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Aligned the license expression with `PEP-639` [#2511](https://github.com/IntelPython/dpnp/pull/2511)
 * Bumped oneMKL version up to `v0.8` [#2514](https://github.com/IntelPython/dpnp/pull/2514)
 * Removed the use of class template argument deduction for alias template to conform to the C++17 standard [#2517](https://github.com/IntelPython/dpnp/pull/2517)
+* Changed th order of individual FFTs over `axes` for `dpnp.fft.irfftn` to be in forward order [#2524](https://github.com/IntelPython/dpnp/pull/2524)
 
 ### Deprecated
 

dpnp/fft/dpnp_utils_fft.py

@@ -110,14 +110,28 @@ def _commit_descriptor(a, forward, in_place, c2c, a_strides, index, batch_fft):
     return dsc, out_strides
 
 
-def _complex_nd_fft(a, s, norm, out, forward, in_place, c2c, axes, batch_fft):
+def _complex_nd_fft(
+    a,
+    s,
+    norm,
+    out,
+    forward,
+    in_place,
+    c2c,
+    axes,
+    batch_fft,
+    *,
+    reversed_axes=True,
+):
     """Computes complex-to-complex FFT of the input N-D array."""
 
     len_axes = len(axes)
     # OneMKL supports up to 3-dimensional FFT on GPU
     # repeated axis in OneMKL FFT is not allowed
     if len_axes > 3 or len(set(axes)) < len_axes:
-        axes_chunk, shape_chunk = _extract_axes_chunk(axes, s, chunk_size=3)
+        axes_chunk, shape_chunk = _extract_axes_chunk(
+            axes, s, chunk_size=3, reversed_axes=reversed_axes
+        )
         for i, (s_chunk, a_chunk) in enumerate(zip(shape_chunk, axes_chunk)):
             a = _truncate_or_pad(a, shape=s_chunk, axes=a_chunk)
             # if out is used in an intermediate step, it will have memory
@@ -291,10 +305,10 @@ def _copy_array(x, complex_input):
     return x, copy_flag
 
 
-def _extract_axes_chunk(a, s, chunk_size=3):
+def _extract_axes_chunk(a, s, chunk_size=3, reversed_axes=True):
     """
     Classify the first input into a list of lists with each list containing
-    only unique values in reverse order and its length is at most `chunk_size`.
+    only unique values and its length is at most `chunk_size`.
     The second input is also classified into a list of lists with each list
     containing the corresponding values of the first input.
 
@@ -306,13 +320,14 @@ def _extract_axes_chunk(a, s, chunk_size=3):
         The second input.
     chunk_size : int
         Maximum number of elements in each chunk.
+    reversed_axes : bool
+        If True, the output chunks will be in reverse order.
 
     Return
     ------
     out : a tuple of two lists
         The first element of output is a list of lists with each list
-        containing only unique values in revere order and its length is
-        at most `chunk_size`.
+        containing only unique values and its length is at most `chunk_size`.
         The second element of output is a list of lists with each list
         containing the corresponding values of the first input.
 
@@ -362,7 +377,10 @@ def _extract_axes_chunk(a, s, chunk_size=3):
         a_chunks.append(a_current_chunk[::-1])
         s_chunks.append(s_current_chunk[::-1])
 
-    return a_chunks[::-1], s_chunks[::-1]
+    if reversed_axes:
+        return a_chunks[::-1], s_chunks[::-1]
+
+    return a_chunks, s_chunks
 
 
 def _fft(a, norm, out, forward, in_place, c2c, axes, batch_fft=True):
@@ -531,9 +549,12 @@ def _validate_out_keyword(a, out, s, axes, c2c, c2r, r2c):
             expected_shape[axes[-1]] = s[-1] // 2 + 1
         elif c2c:
             expected_shape[axes[-1]] = s[-1]
-        for s_i, axis in zip(s[-2::-1], axes[-2::-1]):
-            expected_shape[axis] = s_i
+        if r2c or c2c:
+            for s_i, axis in zip(s[-2::-1], axes[-2::-1]):
+                expected_shape[axis] = s_i
         if c2r:
+            for s_i, axis in zip(s[:-1], axes[:-1]):
+                expected_shape[axis] = s_i
             expected_shape[axes[-1]] = s[-1]
 
         if out.shape != tuple(expected_shape):
@@ -717,6 +738,7 @@ def dpnp_fftn(a, forward, real, s=None, axes=None, norm=None, out=None):
             c2c=True,
             axes=axes[:-1],
             batch_fft=a.ndim != len_axes - 1,
+            reversed_axes=False,
         )
         a = _truncate_or_pad(a, (s[-1],), (axes[-1],))
         if c2r:

dpnp/tests/test_fft.py

@@ -401,13 +401,13 @@ def test_repeated_axes(self, axes):
         result = dpnp.fft.fftn(ia, axes=axes)
         # Intel NumPy ignores repeated axes (mkl_fft-gh-104), handle it one by one
         expected = a
-        for ii in axes:
+        for ii in axes[::-1]:
             expected = numpy.fft.fft(expected, axis=ii)
         assert_dtype_allclose(result, expected)
 
         # inverse FFT
         result = dpnp.fft.ifftn(result, axes=axes)
-        for ii in axes:
+        for ii in axes[::-1]:
             expected = numpy.fft.ifft(expected, axis=ii)
         assert_dtype_allclose(result, expected)
 
@@ -893,7 +893,7 @@ def test_repeated_axes(self, axes):
 
         # inverse FFT
         result = dpnp.fft.irfftn(result, axes=axes)
-        for ii in axes[-2::-1]:
+        for ii in axes[:-1]:
             expected = numpy.fft.ifft(expected, axis=ii)
         expected = numpy.fft.irfft(expected, axis=axes[-1])
         assert_dtype_allclose(result, expected)
@@ -912,7 +912,7 @@ def test_repeated_axes_with_s(self, axes, s):
         assert_dtype_allclose(result, expected)
 
         result = dpnp.fft.irfftn(result, s=s, axes=axes)
-        for jj, ii in zip(s[-2::-1], axes[-2::-1]):
+        for jj, ii in zip(s[:-1], axes[:-1]):
             expected = numpy.fft.ifft(expected, n=jj, axis=ii)
         expected = numpy.fft.irfft(expected, n=s[-1], axis=axes[-1])
         assert_dtype_allclose(result, expected)
@@ -934,7 +934,7 @@ def test_out(self, axes, s):
         assert_dtype_allclose(result, expected)
 
         # inverse FFT
-        for jj, ii in zip(s[-2::-1], axes[-2::-1]):
+        for jj, ii in zip(s[:-1], axes[:-1]):
             expected = numpy.fft.ifft(expected, n=jj, axis=ii)
         expected = numpy.fft.irfft(expected, n=s[-1], axis=axes[-1])
         out = dpnp.empty(expected.shape, dtype=numpy.float32)