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?

[antonwolfy]

Do we need to add copyright header?

[antonwolfy]

Links to PR are missing

[antonwolfy]

Empty axis set → scale 1.0, matches numpy:

Suggested change

	else:
	elif len(axes) == 0:
	return ()
	else:

It probably would be good to cover that in tests

[antonwolfy]

Nice fix — the norm×subset-axes and c2r output-length bugs are correctly resolved, and the new test_dispatch_equivalence.py fails loudly on master (74 failures) and passes cleanly after. Two optional follow-ups on test coverage; I verified every case below passes against the fixed backend, so these are regression guards for untested-but-correct paths, not new bugs.

Gap A — out= coverage

The only out= test (test_fftn_axes_subset_out) covers fftn / complex only / no norm.

Missing case	Why it matters	Example
out= + scaled norm=	The PR's whole subject — the scale is written into out; nothing checks the scaled-write path	mkl_fft.fftn(x, axes=(0,), norm="forward", out=out)
out= on ifftn	Shares the out path but never exercised with it	mkl_fft.ifftn(x, out=out)
out= on rfftn	r2c: out is complex with reduced last axis n//2+1 — different allocation than c2c	out = np.empty((8,7,7), complex128); mkl_fft.rfftn(xr, out=out)
out= on irfftn (+ norm)	c2r: out is real with the expanded 2*(n-1) last axis — exactly the length the invreal branch computes	out = np.empty((8,7,24), float64); mkl_fft.irfftn(xc, norm="forward", out=out)

Gap B — s= combined with a scaled norm

The helper deliberately early-returns when s is not None, so _compute_fwd_scale uses prod(s). That branch is only tested without norm; the s= × scaled-norm intersection is untested.

Missing case	Why it matters	Example
s= pads a subset axis + forward/ortho	Scale must be prod(s) (padded length), not the original axis length	mkl_fft.fftn(x, s=(16,), axes=(0,), norm="forward")
s= truncates + scaled norm	Same, downward	mkl_fft.fftn(x, s=(4,), axes=(0,), norm="ortho")
irfftn with explicit s + scaled norm	When s is given, the invreal doubling 2*(n-1) must not apply — a regression that double-applies it only surfaces here	mkl_fft.irfftn(xc, s=(8,7,20), norm="forward") → normalizes over 20, not 2*(13-1)
rfftn with s= on a subset + scaled norm	Locks the r2c s-given path with scaling	mkl_fft.rfftn(xr, s=(10,), axes=(1,), norm="ortho")