Distributed Fast Fourier Transforms

heat/__init__.py

@@ -9,6 +9,7 @@
 from . import core
 from . import classification
 from . import cluster
+from . import fft
 from . import graph
 from . import naive_bayes
 from . import nn

heat/core/stride_tricks.py

@@ -151,7 +151,10 @@ def sanitize_axis(
 
     """
     # scalars are handled like unsplit matrices
-    if len(shape) == 0:
+    original_axis = axis
+    ndim = len(shape)
+
+    if ndim == 0:
         axis = None
 
     if axis is not None and not isinstance(axis, int) and not isinstance(axis, tuple):
@@ -160,7 +163,9 @@ def sanitize_axis(
         axis = tuple(dim + len(shape) if dim < 0 else dim for dim in axis)
         for dim in axis:
             if dim < 0 or dim >= len(shape):
-                raise ValueError(f"axis {axis} is out of bounds for shape {shape}")
+                raise ValueError(
+                    f"axis {original_axis} is out of bounds for {ndim}-dimensional array"
+                )
         return axis
 
     if axis is None or 0 <= axis < len(shape):
@@ -169,7 +174,7 @@ def sanitize_axis(
         axis += len(shape)
 
     if axis < 0 or axis >= len(shape):
-        raise ValueError(f"axis {axis} is out of bounds for shape {shape}")
+        raise ValueError(f"axis {original_axis} is out of bounds for {ndim}-dimensional array")
 
     return axis
 

heat/core/tests/test_suites/basic_test.py

@@ -129,11 +129,11 @@ def assert_array_equal(self, heat_array, expected_array):
             f"Local shapes do not match. Got {heat_array.lshape} expected {expected_array[slices].shape}",
         )
         # compare local tensors to corresponding slice of expected_array
-        is_allclose = np.allclose(heat_array.larray.cpu(), expected_array[slices])
-        ht_is_allclose = ht.array(
-            [is_allclose], dtype=ht.bool, is_split=0, device=heat_array.device
+        is_allclose = torch.tensor(
+            np.allclose(heat_array.larray.cpu(), expected_array[slices]), dtype=torch.int32
         )
-        self.assertTrue(ht.all(ht_is_allclose))
+        heat_array.comm.Allreduce(MPI.IN_PLACE, is_allclose, MPI.SUM)
+        self.assertTrue(is_allclose == heat_array.comm.size)
 
     def assert_func_equal(
         self,

heat/fft/__init__.py

@@ -0,0 +1,5 @@
+"""
+import the fft functions into the fft namespace
+"""
+
+from .fft import *