address comemnts

Author: Vahid Tavanashad

mkl_fft/_numpy_fft.py

@@ -75,7 +75,7 @@ def frwd_sc_1d(n, s):
     return 1/nn if nn != 0 else 1
 
 
-def frwd_sc_nd(s, axes, x_shape):
+def frwd_sc_nd(s, x_shape):
     ss = s if s is not None else x_shape
     nn = prod(ss)
     return 1/nn if nn != 0 else 1
@@ -812,9 +812,9 @@ def fftn(a, s=None, axes=None, norm=None):
     if norm in (None, "backward"):
         fsc = 1.0
     elif norm == "forward":
-        fsc = frwd_sc_nd(s, axes, x.shape)
+        fsc = frwd_sc_nd(s, x.shape)
     else:
-        fsc = sqrt(frwd_sc_nd(s, axes, x.shape))
+        fsc = sqrt(frwd_sc_nd(s, x.shape))
 
     return trycall(
         mkl_fft.fftn,
@@ -928,9 +928,9 @@ def ifftn(a, s=None, axes=None, norm=None):
     if norm in (None, "backward"):
         fsc = 1.0
     elif norm == "forward":
-        fsc = frwd_sc_nd(s, axes, x.shape)
+        fsc = frwd_sc_nd(s, x.shape)
     else:
-        fsc = sqrt(frwd_sc_nd(s, axes, x.shape))
+        fsc = sqrt(frwd_sc_nd(s, x.shape))
 
     return trycall(
         mkl_fft.ifftn,
@@ -1227,11 +1227,11 @@ def rfftn(a, s=None, axes=None, norm=None):
     elif norm == "forward":
         x = asanyarray(x)
         s, axes = _cook_nd_args(x, s, axes)
-        fsc = frwd_sc_nd(s, axes, x.shape)
+        fsc = frwd_sc_nd(s, x.shape)
     else:
         x = asanyarray(x)
         s, axes = _cook_nd_args(x, s, axes)
-        fsc = sqrt(frwd_sc_nd(s, axes, x.shape))
+        fsc = sqrt(frwd_sc_nd(s, x.shape))
 
     return trycall(
         mkl_fft.rfftn,
@@ -1384,11 +1384,11 @@ def irfftn(a, s=None, axes=None, norm=None):
     elif norm == "forward":
         x = asanyarray(x)
         s, axes = _cook_nd_args(x, s, axes, invreal=1)
-        fsc = frwd_sc_nd(s, axes, x.shape)
+        fsc = frwd_sc_nd(s, x.shape)
     else:
         x = asanyarray(x)
         s, axes = _cook_nd_args(x, s, axes, invreal=1)
-        fsc = sqrt(frwd_sc_nd(s, axes, x.shape))
+        fsc = sqrt(frwd_sc_nd(s, x.shape))
 
     return trycall(
         mkl_fft.irfftn,

mkl_fft/_scipy_fft.py

@@ -200,16 +200,13 @@ def _check_plan(plan):
 
 
 def _frwd_sc_1d(n, s):
-    nn = n if n else s
+    nn = n if n is not None else s
     return 1/nn if nn != 0 else 1
 
 
-def _frwd_sc_nd(s, axes, x_shape):
+def _frwd_sc_nd(s, x_shape):
     ss = s if s is not None else x_shape
-    if axes is not None:
-        nn = prod([ss[ai] for ai in axes])
-    else:
-        nn = prod(ss)
+    nn = prod(ss)
     return 1/nn if nn != 0 else 1
 
 
@@ -233,9 +230,9 @@ def _compute_nd_fwd_scale(norm, s, axes, x_shape):
     if norm in (None, "backward"):
         fsc = 1.0
     elif norm == "forward":
-        fsc = _frwd_sc_nd(s, axes, x_shape)
+        fsc = _frwd_sc_nd(s, x_shape)
     elif norm == "ortho":
-        fsc = sqrt(_frwd_sc_nd(s, axes, x_shape))
+        fsc = sqrt(_frwd_sc_nd(s, x_shape))
     else:
         _check_norm(norm)
     return fsc
@@ -359,10 +356,10 @@ def _compute_nd_fwd_scale_for_rfft(norm, s, axes, x, invreal=False):
         fsc = 1.0
     elif norm == "forward":
         s, axes = _cook_nd_args(x, s, axes, invreal=invreal)
-        fsc = _frwd_sc_nd(s, axes, x.shape)
+        fsc = _frwd_sc_nd(s, x.shape)
     elif norm == "ortho":
         s, axes = _cook_nd_args(x, s, axes, invreal=invreal)
-        fsc = sqrt(_frwd_sc_nd(s, axes, x.shape))
+        fsc = sqrt(_frwd_sc_nd(s, x.shape))
     else:
         _check_norm(norm)
     return s, axes, fsc