fix: norm scaling

Author: jharlow-intel

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"` and `norm="ortho"` scaling in `mkl_fft.fftn`, `ifftn`, `rfftn`, `irfftn` and the `fft2`/`ifft2`/`rfft2`/`irfft2` family when only a subset of the input axes is transformed and `s` is not given. The scale factor was computed over the full array shape instead of over the transformed axes, over-normalizing the result by the product of the untransformed axis lengths (for example `fft2` on a 3-D array, or `fftn(x, axes=(0,))`). The `mkl_fft.interfaces.numpy_fft` and `mkl_fft.interfaces.scipy_fft` wrappers were unaffected, as they resolve `s` before delegating
+* Fixed `norm="forward"` and `norm="ortho"` scaling in `mkl_fft.irfftn` and `irfft2`, which normalized over the input length `n` along the last transformed axis rather than the complex-to-real output length `2 * (n - 1)`. This applied even when every axis was transformed
 * 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,43 @@ def _compute_fwd_scale(norm, n, shape):
         return np.sqrt(fsc)
 
 
+def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False):
+    """
+    Resolve the lengths that a norm-scaled N-D transform normalizes over.
+
+    ``_compute_fwd_scale`` falls back to the full array shape when ``s`` is
+    None. That over-normalizes when only a subset of axes is transformed, and
+    for c2r transforms the basis is the *output* length along the last
+    transformed axis rather than the input one.
+
+    This mirrors what the ``numpy_fft`` and ``scipy_fft`` interfaces already
+    do by calling ``_cook_nd_args`` before delegating. Only the scale basis is
+    resolved here; ``s`` itself is left alone so that dispatch in
+    ``_c2c_fftnd_impl`` is unchanged.
+
+    ``norm`` is accepted only to skip the work for the unscaled norms, whose
+    scale is 1.0 regardless of shape. Invalid values fall through to
+    ``_compute_fwd_scale``, which validates them.
+    """
+
+    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
+        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 or empty axes; leave the scale alone and let the
+        # transform itself raise
+        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)

mkl_fft/tests/test_dispatch_equivalence.py

@@ -0,0 +1,230 @@
+"""Cross-library equivalence checks for axis and axes dispatch.
+
+``third_party/scipy/test_basic.py::test_fft_with_order`` already checks that
+mkl_fft agrees with *itself* across C, Fortran, and non-contiguous layouts. It
+does not compare against an external reference, so a dispatch change that is
+consistently wrong in every layout passes it.
+
+The defect recorded in ``_fft_utils._iter_complementary`` was exactly that
+kind: values correct, but an element placed in the other half of the output
+relative to NumPy. These tests therefore use ``numpy.fft`` as the reference.
+
+Two deliberate choices:
+
+* Every axis length differs, so an axis permutation cannot produce a
+  correctly shaped result and hide behind a shape assertion.
+* Output dtype is asserted alongside values, so a dispatch change cannot
+  silently upcast.
+
+These cover the paths that dispatch on *which* axes are requested: full-axes
+transforms reach the batched N-D descriptor, strict subsets iterate the
+complementary axes, and 1-D transforms of rank > 2 arrays are batched only for
+the first and last axis.
+"""
+
+import itertools
+
+import numpy as np
+import pytest
+from numpy.testing import assert_allclose
+
+import mkl_fft
+
+_SHAPE_3D = (8, 7, 13)
+_SHAPE_4D = (4, 5, 6, 7)
+
+_DTYPES = ["float32", "float64", "complex64", "complex128"]
+_REAL_DTYPES = ["float32", "float64"]
+
+_ORDERS = ["C", "F", "non-contiguous"]
+
+# Relative tolerance by input precision. Single-precision transforms of
+# random data over these lengths stay well inside 2e-5.
+_TOL = {
+    "float32": 2e-5,
+    "complex64": 2e-5,
+    "float64": 1e-12,
+    "complex128": 1e-12,
+}
+
+# every non-empty subset of the axes of a 3-D array, plus None
+_AXES_3D = [
+    ax for n in (1, 2, 3) for ax in itertools.combinations(range(3), n)
+] + [None]
+
+
+def _make(shape, dtype, seed=42):
+    rng = np.random.default_rng(seed)
+    dt = np.dtype(dtype)
+    if dt.kind == "c":
+        x = rng.standard_normal(shape) + 1j * rng.standard_normal(shape)
+    else:
+        x = rng.standard_normal(shape)
+    return x.astype(dt)
+
+
+def _relayout(x, order):
+    """Return *x* laid out as requested; data content may differ by order."""
+    if order == "F":
+        return np.asfortranarray(x)
+    if order == "non-contiguous":
+        return x[::-1]
+    return np.ascontiguousarray(x)
+
+
+def _check(got, want, dtype):
+    assert got.dtype == want.dtype, f"dtype {got.dtype} != {want.dtype}"
+    assert got.shape == want.shape, f"shape {got.shape} != {want.shape}"
+    tol = _TOL[dtype]
+    assert_allclose(
+        got, want, rtol=tol, atol=tol * max(1.0, float(np.abs(want).max()))
+    )
+
+
+# ---------------------------------------------------------------------------
+# N-D complex transforms over a subset of axes
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("func", ["fftn", "ifftn"])
+@pytest.mark.parametrize("dtype", _DTYPES)
+@pytest.mark.parametrize("axes", _AXES_3D)
+@pytest.mark.parametrize("order", _ORDERS)
+def test_fftn_axes_subset(func, dtype, axes, order):
+    x = _relayout(_make(_SHAPE_3D, dtype), order)
+    got = getattr(mkl_fft, func)(x, axes=axes)
+    want = getattr(np.fft, func)(x, axes=axes)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["rfftn", "irfftn"])
+@pytest.mark.parametrize("dtype", _DTYPES)
+@pytest.mark.parametrize("axes", _AXES_3D)
+@pytest.mark.parametrize("order", _ORDERS)
+def test_rfftn_axes_subset(func, dtype, axes, order):
+    if func == "rfftn" and dtype not in _REAL_DTYPES:
+        pytest.skip("rfftn takes real input")
+    x = _relayout(_make(_SHAPE_3D, dtype), order)
+    got = getattr(mkl_fft, func)(x, axes=axes)
+    want = getattr(np.fft, func)(x, axes=axes)
+    _check(got, want, dtype)
+
+
+# ---------------------------------------------------------------------------
+# 1-D transforms along each axis of a higher-rank array
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("func", ["fft", "ifft"])
+@pytest.mark.parametrize("dtype", _DTYPES)
+@pytest.mark.parametrize("axis", range(len(_SHAPE_3D)))
+@pytest.mark.parametrize("order", _ORDERS)
+def test_fft_axis_3d(func, dtype, axis, order):
+    x = _relayout(_make(_SHAPE_3D, dtype), order)
+    got = getattr(mkl_fft, func)(x, axis=axis)
+    want = getattr(np.fft, func)(x, axis=axis)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["fft", "ifft", "rfft"])
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("axis", range(len(_SHAPE_4D)))
+@pytest.mark.parametrize("order", _ORDERS)
+def test_fft_axis_4d(func, dtype, axis, order):
+    """A rank-4 array has two interior axes, so the per-vector fallback in the
+    C backend is exercised twice within one sweep.
+    """
+    if func == "rfft" and dtype != "float64":
+        pytest.skip("rfft takes real input")
+    x = _relayout(_make(_SHAPE_4D, dtype), order)
+    got = getattr(mkl_fft, func)(x, axis=axis)
+    want = getattr(np.fft, func)(x, axis=axis)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["rfft", "irfft"])
+@pytest.mark.parametrize("dtype", _DTYPES)
+@pytest.mark.parametrize("axis", range(len(_SHAPE_3D)))
+@pytest.mark.parametrize("order", _ORDERS)
+def test_rfft_axis_3d(func, dtype, axis, order):
+    if func == "rfft" and dtype not in _REAL_DTYPES:
+        pytest.skip("rfft takes real input")
+    x = _relayout(_make(_SHAPE_3D, dtype), order)
+    got = getattr(mkl_fft, func)(x, axis=axis)
+    want = getattr(np.fft, func)(x, axis=axis)
+    _check(got, want, dtype)
+
+
+# ---------------------------------------------------------------------------
+# norm interacts with the scale factor applied at dispatch time
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("func", ["fftn", "ifftn"])
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("axes", [(0,), (1,), (2,), (1, 2), None])
+@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
+def test_fftn_axes_subset_norm(func, dtype, axes, norm):
+    x = _make(_SHAPE_3D, dtype)
+    got = getattr(mkl_fft, func)(x, axes=axes, norm=norm)
+    want = getattr(np.fft, func)(x, axes=axes, norm=norm)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["rfftn", "irfftn"])
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("axes", [(0,), (1,), (2,), (1, 2), None])
+@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
+def test_rfftn_axes_subset_norm(func, dtype, axes, norm):
+    """Includes ``axes=None``: for c2r the scale basis is the *output* length
+    along the last transformed axis, so a full-axes irfftn is normalized over
+    ``2 * (n - 1)`` rather than ``n``.
+    """
+    if func == "rfftn" and dtype != "float64":
+        pytest.skip("rfftn takes real input")
+    x = _make(_SHAPE_3D, dtype)
+    got = getattr(mkl_fft, func)(x, axes=axes, norm=norm)
+    want = getattr(np.fft, func)(x, axes=axes, norm=norm)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["fft2", "ifft2", "rfft2", "irfft2"])
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
+def test_fft2_on_3d_norm(func, dtype, norm):
+    """``fft2`` on a rank-3 array transforms 2 of 3 axes, so it is a subset
+    transform even though the caller passed no ``axes``.
+    """
+    if func == "rfft2" and dtype != "float64":
+        pytest.skip("rfft2 takes real input")
+    x = _make(_SHAPE_3D, dtype)
+    got = getattr(mkl_fft, func)(x, norm=norm)
+    want = getattr(np.fft, func)(x, norm=norm)
+    _check(got, want, dtype)
+
+
+@pytest.mark.parametrize("func", ["fft", "ifft"])
+@pytest.mark.parametrize("dtype", ["float64", "complex128"])
+@pytest.mark.parametrize("axis", range(len(_SHAPE_3D)))
+@pytest.mark.parametrize("norm", [None, "backward", "forward", "ortho"])
+def test_fft_axis_norm(func, dtype, axis, norm):
+    x = _make(_SHAPE_3D, dtype)
+    got = getattr(mkl_fft, func)(x, axis=axis, norm=norm)
+    want = getattr(np.fft, func)(x, axis=axis, norm=norm)
+    _check(got, want, dtype)
+
+
+# ---------------------------------------------------------------------------
+# out= must not change results on any dispatch path
+# ---------------------------------------------------------------------------
+
+
+@pytest.mark.parametrize("dtype", ["complex64", "complex128"])
+@pytest.mark.parametrize("axes", _AXES_3D)
+def test_fftn_axes_subset_out(dtype, axes):
+    x = _make(_SHAPE_3D, dtype)
+    want = np.fft.fftn(x, axes=axes)
+    out = np.empty(want.shape, dtype=x.dtype)
+    got = mkl_fft.fftn(x, axes=axes, out=out)
+    assert got is out, "out= should be returned"
+    _check(got, want, dtype)