Merge pull request #111 from IntelPython/fix-gh-109

Fix gh-109
IntelPython · Jul 19, 2024 · c13a9a2 · c13a9a2
2 parents 6335c59 + abffaec
commit c13a9a2
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 8 deletions.
diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml
@@ -41,7 +41,7 @@ jobs:
         run: conda install conda-build
       - name: Build conda package
         run: |
-          CHANNELS="-c conda-forge -c intel --override-channels"
+          CHANNELS="-c conda-forge -c https://software.repos.intel.com/python/conda --override-channels"
           VERSIONS="--python ${{ matrix.python }}"
           TEST="--no-test"
 
@@ -67,7 +67,7 @@ jobs:
         runner: [ubuntu-latest]
     continue-on-error: ${{ matrix.experimental }}
     env:
-      CHANNELS: -c intel -c main --override-channels
+      CHANNELS: -c https://software.repos.intel.com/python/conda -c main --override-channels
 
     steps:
       - name: Download artifact
@@ -150,7 +150,7 @@ jobs:
             ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-python-${{ matrix.python }}-
             ${{ runner.os }}-conda-${{ env.CACHE_NUMBER }}-
       - name: Build conda package
-        run: conda build --no-test --python ${{ matrix.python }} -c intel -c conda-forge --override-channels conda-recipe
+        run: conda build --no-test --python ${{ matrix.python }} -c https://software.repos.intel.com/python/conda -c conda-forge --override-channels conda-recipe
       - name: Upload artifact
         uses: actions/upload-artifact@v4
         with:
@@ -171,7 +171,7 @@ jobs:
     continue-on-error: ${{ matrix.experimental }}
     env:
       workdir: '${{ github.workspace }}'
-      CHANNELS: -c intel -c conda-forge --override-channels
+      CHANNELS: -c https://software.repos.intel.com/python/conda -c conda-forge --override-channels
 
     steps:
       - name: Download artifact

diff --git a/mkl_fft/_pydfti.pyx b/mkl_fft/_pydfti.pyx
@@ -901,12 +901,12 @@ def _cook_nd_args(a, s=None, axes=None, invreal=0):
     return s, axes
 
 
-def _iter_fftnd(a, s=None, axes=None, function=fft, overwrite_arg=False, scale_function=lambda n: 1.0):
+def _iter_fftnd(a, s=None, axes=None, function=fft, overwrite_arg=False, scale_function=lambda n, ind: 1.0):
     a = np.asarray(a)
     s, axes = _init_nd_shape_and_axes(a, s, axes)
     ovwr = overwrite_arg
     for ii in reversed(range(len(axes))):
-        a = function(a, n = s[ii], axis = axes[ii], overwrite_x=ovwr, forward_scale=scale_function(s[ii]))
+        a = function(a, n = s[ii], axis = axes[ii], overwrite_x=ovwr, forward_scale=scale_function(s[ii], ii))
         ovwr = True
     return a
 
@@ -1093,9 +1093,9 @@ def _fftnd_impl(x, shape=None, axes=None, overwrite_x=False, direction=+1, doubl
                 res
                 )
         else:
-            sc = (<object> fsc)**(1/x.ndim)
+            sc = <object> fsc
             return _iter_fftnd(x, s=shape, axes=axes,
-                               overwrite_arg=overwrite_x, scale_function=lambda n: sc,
+                               overwrite_arg=overwrite_x, scale_function=lambda n, i: sc if i == 0 else 1.,
                                function=fft if direction == 1 else ifft)
 
 

diff --git a/mkl_fft/tests/test_fftnd.py b/mkl_fft/tests/test_fftnd.py
@@ -217,3 +217,14 @@ def test_scale_nd_axes(self):
 
         r_tol, a_tol = _get_rtol_atol(X)
         assert_allclose(f, 5*f_scale, rtol=r_tol, atol=a_tol)
+
+
+def test_gh109():
+    b_int = np.array([[5, 7, 6, 5], [4, 6, 4, 8], [9, 3, 7, 5]], dtype=np.int64)
+    b = np.asarray(b_int, dtype=np.float32)
+
+    r1 = mkl_fft.fftn(b, shape=None, axes=(0,), overwrite_x=False, forward_scale=1/3)
+    r2 = mkl_fft.fftn(b_int, shape=None, axes=(0,), overwrite_x=False, forward_scale=1/3)
+
+    rtol, atol = _get_rtol_atol(b)
+    assert_allclose(r1, r2, rtol=rtol, atol=atol)