fix: review

Author: jharlow-intel

mkl_fft/_fft_utils.py

@@ -80,21 +80,13 @@ def _compute_fwd_scale(norm, n, shape):
 
 def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False):
     """
-    Resolve the lengths that a norm-scaled N-D transform normalizes over.
+    Lengths 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.
+    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"):
@@ -104,17 +96,15 @@ def _compute_nd_scale_shape(x, s, axes, norm=None, invreal=False):
             ss = list(x.shape)
             last = len(ss) - 1
         elif len(axes) == 0:
-            # transforming no axes is an identity, so the scale is 1.0;
-            # np.prod(()) is 1, which gives that for every norm
+            # 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; leave the scale alone and let the transform itself
-        # raise, so the error matches what NumPy reports
+        # invalid axes; let the transform raise, matching NumPy's error
         return s
     return tuple(ss)