fix: norm scaling

[jharlow-intel]

Doing some iterative agentic looping dev experimentation. Part of it found this:

Two norm scaling bugs in the root N-D API returned silently mis-scaled
results. Both are fixed by resolving the scale basis from the transformed
axes, matching what mkl_fft.interfaces.* already does via _cook_nd_args.

1. Subset-axes transforms over-normalized by the product of the
   untransformed axis lengths, because the scale was computed over the full
   array shape. Hits fftn(x, axes=(0,)), and less obviously fft2(x) on a
   3-D array — that transforms 2 of 3 axes.
2. irfftn/irfft2 normalized over the input length n rather than the
   complex-to-real output length 2 * (n - 1) along the last transformed
   axis. Wrong even when every axis was transformed.

Applies to fftn/ifftn/rfftn/irfftn and the fft2/ifft2/rfft2/
irfft2 family with norm="forward" or "ortho" and no explicit s.
Unaffected: norm=None/"backward", explicit s=, 1-D transforms, and
interfaces.numpy_fft/scipy_fft.

Why it wasn't caught

The existing N-D norm tests compare mkl_fft against other mkl_fft calls,
and test_fft_with_order compares it against itself across memory layouts —
self-consistency, never an external reference. The new
test_dispatch_equivalence.py uses numpy.fft as the reference across
dtype × layout × axes × norm, on a shape whose axis lengths all differ so that
an axis permutation cannot produce a correctly shaped result.

Testing

1725 passed / 104 skipped (existing suite was 971 — no regressions).
176 root-API combinations checked against numpy.fft: 0 mismatches, 32 of
them failing before the fix. The norm=None path is unchanged; the new helper
short-circuits in ~0.04 µs.

This is part of other on-going performance improvement looping I'm doing. No rush in merging it, it was entirely agentic and I didn't have time to thoroughly review, hopefully an expert here can say whether the PR is correct or not

[Copilot]

Pull request overview

Fixes incorrect norm="forward" / "ortho" scaling in the root N-D FFT API (mkl_fft.fftn/ifftn/rfftn/irfftn and, via delegation, the fft2/ifft2/rfft2/irfft2 family) by computing the normalization basis from the transformed axes (and for irfftn from the complex-to-real output length along the last transformed axis), aligning behavior with the existing interface wrappers.

Changes:

• Introduced _compute_nd_scale_shape(...) to derive the correct scale basis for N-D transforms when norm is scaled and s is not provided.
• Updated root N-D entry points to use the derived scale basis when computing fsc.
• Added a comprehensive NumPy-reference equivalence test suite covering dtype × layout × axes × norm dispatch paths; documented the fix in the changelog.

Reviewed changes

Copilot reviewed 4 out of 4 changed files in this pull request and generated no comments.

File	Description
mkl_fft/_fft_utils.py	Adds _compute_nd_scale_shape to compute the correct normalization basis for scaled norms in N-D transforms (including irfftn output-length handling).
mkl_fft/_mkl_fft.py	Switches root N-D FFT wrappers to compute fsc from the transformed-axis scale basis instead of the full array shape.
mkl_fft/tests/test_dispatch_equivalence.py	Adds NumPy-reference dispatch/equivalence tests to catch axis/axes/norm scaling and dispatch regressions.
CHANGELOG.md	Documents the scaling fixes for subset-axes transforms and irfftn/irfft2 output-length normalization.

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[jharlow-intel]

reproducer:

conda create -n mkl_fft_test --override-channels -c https://software.repos.intel.com/python/conda -c conda-forge python mkl_fft

conda activate mkl_fft_test

# python execute following script

import sys

import numpy as np

import mkl_fft

x = np.random.default_rng(0).standard_normal((8, 7, 13)) + 0j
bad = 0


def check(label, got, want):
    global bad
    f = np.vdot(want, got) / np.vdot(want, want)  # least-squares scale
    pure = np.allclose(got, f * want)  # wrong by ONLY that scale?
    ok = abs(f - 1) < 1e-9 and pure
    bad += not ok
    note = "" if pure else "   <-- not a pure scale, values differ too!"
    print(f"  {'ok ' if ok else 'BUG'}  {label:<38} scale={f.real:9.6f}{note}")


print(f"mkl_fft {mkl_fft.__version__}, numpy {np.__version__}, x.shape={x.shape}\n")

print("subset of axes, s not given:")
for axes in [(0,), (1,), (2,), (1, 2)]:
    for norm in ("forward", "ortho"):
        check(
            f"fftn(axes={axes}, norm={norm!r})",
            mkl_fft.fftn(x, axes=axes, norm=norm),
            np.fft.fftn(x, axes=axes, norm=norm),
        )

print("fft2 on a 3-D array -- transforms 2 of 3 axes:")
for norm in ("forward", "ortho"):
    check(
        f"fft2(norm={norm!r})",
        mkl_fft.fft2(x, norm=norm),
        np.fft.fft2(x, norm=norm),
    )

print("complex-to-real, every axis transformed:")
for fn in ("irfftn", "irfft2"):
    for norm in ("forward", "ortho"):
        check(
            f"{fn}(norm={norm!r})",
            getattr(mkl_fft, fn)(x, norm=norm),
            getattr(np.fft, fn)(x, norm=norm),
        )

print("\ncontrols that should always pass:")
check("fftn(axes=None, norm='ortho')", mkl_fft.fftn(x, norm="ortho"), np.fft.fftn(x, norm="ortho"))
check("fft(axis=1, norm='ortho')", mkl_fft.fft(x, axis=1, norm="ortho"), np.fft.fft(x, axis=1, norm="ortho"))
check("fftn(axes=(0,), norm=None)", mkl_fft.fftn(x, axes=(0,)), np.fft.fftn(x, axes=(0,)))

print(f"\n{bad} mismatched -> bug present" if bad else "\nall match -> fixed")
sys.exit(1 if bad else 0)

[ndgrigorian]

@jharlow-intel can you check if this also covers #336?

[jharlow-intel]

@ndgrigorian this seems to indeed fix the issue detailed in #336

[jharlow-intel]

@ndgrigorian @vlad-perevezentsev added some tests to cover what that user ran into. This is probably ready for a review from the experts
CI failures were due to internal channel issue, they should work on next run