fix: norm scaling

CHANGELOG.md

@@ -11,6 +11,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Changed
 
 ### Fixed
+* Fixed `norm="forward"`/`"ortho"` scaling in `fftn`, `ifftn`, `rfftn`, `irfftn` and the `fft2` family when only a subset of axes is transformed: the scale used the full array shape instead of the transformed axes. The `numpy_fft` and `scipy_fft` interfaces were unaffected [gh-370](https://github.com/IntelPython/mkl_fft/pull/370)
+* Fixed `norm="forward"`/`"ortho"` scaling in `irfftn` and `irfft2`, which normalized over the input length `n` rather than the complex-to-real output length `2 * (n - 1)` [gh-370](https://github.com/IntelPython/mkl_fft/pull/370)
 * Declared `f_ndim` as a C `int` in `_allocate_result` so the buffer size is computed in C rather than through a Python object, resolving a Coverity out-of-bounds (OVERRUN) false positive [gh-364](https://github.com/IntelPython/mkl_fft/pull/364)
 * Silenced a Coverity `UNUSED_VALUE` finding in `__create_descriptor_1d` by marking the `DftiFreeDescriptor` status (used only by a debug-only `assert`) as intentionally unused [gh-365](https://github.com/IntelPython/mkl_fft/pull/365)
 

mkl_fft/_fft_utils.py

@@ -78,6 +78,37 @@ def _compute_fwd_scale(norm, n, shape):
         return np.sqrt(fsc)
 
 
+def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False):
+    """
+    Lengths a norm-scaled N-D transform normalizes over.
+
+    ``_compute_fwd_scale`` falls back to the full array shape when ``s`` is
+    None, which over-normalizes a subset-of-axes transform; for c2r the basis
+    is the output length ``2 * (n - 1)``. Mirrors what the interfaces already
+    do via ``_cook_nd_args``, but leaves ``s`` alone so dispatch is unchanged.
+    ``norm`` is only used to skip the work when the scale is 1.0 anyway.
+    """
+
+    if s is not None or norm in (None, "backward"):
+        return s
+    try:
+        if axes is None:
+            ss = list(x.shape)
+            last = len(ss) - 1
+        elif len(axes) == 0:
+            # identity transform; np.prod(()) == 1 gives scale 1.0
+            return ()
+        else:
+            ss = [x.shape[ai] for ai in axes]
+            last = axes[-1]
+        if invreal:
+            ss[-1] = 2 * (x.shape[last] - 1)
+    except (IndexError, TypeError):
+        # invalid axes; let the transform raise, matching NumPy's error
+        return s
+    return tuple(ss)
+
+
 def _cook_nd_args(a, s=None, axes=None, invreal=False):
     if s is None:
         shapeless = True

mkl_fft/_mkl_fft.py

@@ -27,6 +27,7 @@
     _c2c_fftnd_impl,
     _c2r_fftnd_impl,
     _compute_fwd_scale,
+    _compute_nd_scale_shape,
     _r2c_fftnd_impl,
 )
 
@@ -68,12 +69,14 @@ def ifft2(x, s=None, axes=(-2, -1), norm=None, out=None):
 
 
 def fftn(x, s=None, axes=None, norm=None, out=None):
-    fsc = _compute_fwd_scale(norm, s, x.shape)
+    ss = _compute_nd_scale_shape(x, s, axes, norm)
+    fsc = _compute_fwd_scale(norm, ss, x.shape)
     return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=+1, fsc=fsc)
 
 
 def ifftn(x, s=None, axes=None, norm=None, out=None):
-    fsc = _compute_fwd_scale(norm, s, x.shape)
+    ss = _compute_nd_scale_shape(x, s, axes, norm)
+    fsc = _compute_fwd_scale(norm, ss, x.shape)
     return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=-1, fsc=fsc)
 
 
@@ -96,10 +99,12 @@ def irfft2(x, s=None, axes=(-2, -1), norm=None, out=None):
 
 
 def rfftn(x, s=None, axes=None, norm=None, out=None):
-    fsc = _compute_fwd_scale(norm, s, x.shape)
+    ss = _compute_nd_scale_shape(x, s, axes, norm)
+    fsc = _compute_fwd_scale(norm, ss, x.shape)
     return _r2c_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc)
 
 
 def irfftn(x, s=None, axes=None, norm=None, out=None):
-    fsc = _compute_fwd_scale(norm, s, x.shape)
+    ss = _compute_nd_scale_shape(x, s, axes, norm, invreal=True)
+    fsc = _compute_fwd_scale(norm, ss, x.shape)
     return _c2r_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc)