fix: norm scaling

Author: jharlow-intel

CHANGELOG.md

@@ -14,6 +14,8 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 * Pinned Cython in the Coverity Scan workflow so generated code stays stable between scans, and added `coverity/README.md` documenting the known false-positive families and the scan review checklist [gh-374](https://github.com/IntelPython/mkl_fft/pull/374)
 
 ### 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 [gh-336](https://github.com/IntelPython/mkl_fft/issues/336), [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)
 * Fixed possible memory leaks when `PyMem_Malloc` fails, and raise `MemoryError` [gh-373](https://github.com/IntelPython/mkl_fft/pull/373)

mkl_fft/_fft_utils.py

@@ -102,6 +102,22 @@ def _cook_nd_args(a, s=None, axes=None, invreal=False):
     return s, axes
 
 
+def _nd_fwd_scale(x, s, axes, norm, invreal=False):
+    """Forward scale for an N-D transform, over the transformed axes only.
+
+    ``_compute_fwd_scale`` falls back to the full array shape when ``s`` is
+    None, which over-normalizes a subset of axes; ``_cook_nd_args`` already
+    resolves the right lengths, including c2r's ``2 * (n - 1)``.
+    """
+
+    _check_norm(norm)
+    if norm in (None, "backward"):
+        return 1.0
+    if s is None:
+        s, _ = _cook_nd_args(x, s, axes, invreal=invreal)
+    return _compute_fwd_scale(norm, s, x.shape)
+
+
 # copied from scipy.fft module
 # https://github.com/scipy/scipy/blob/main/scipy/fft/_pocketfft/helper.py
 def _datacopied(arr, original):

mkl_fft/_mkl_fft.py

@@ -27,6 +27,7 @@
     _c2c_fftnd_impl,
     _c2r_fftnd_impl,
     _compute_fwd_scale,
+    _nd_fwd_scale,
     _r2c_fftnd_impl,
 )
 
@@ -68,12 +69,12 @@ 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)
+    fsc = _nd_fwd_scale(x, s, axes, norm)
     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)
+    fsc = _nd_fwd_scale(x, s, axes, norm)
     return _c2c_fftnd_impl(x, s=s, axes=axes, out=out, direction=-1, fsc=fsc)
 
 
@@ -96,10 +97,10 @@ 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)
+    fsc = _nd_fwd_scale(x, s, axes, norm)
     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)
+    fsc = _nd_fwd_scale(x, s, axes, norm, invreal=True)
     return _c2r_fftnd_impl(x, s=s, axes=axes, out=out, fsc=fsc)