Distributed Fast Fourier Transforms (#1218)

* implement fftn, first draft

* implement general , add

* split fft_op and fftn_op, implement inverse and real fft

* add TODO hermitian fft

* add fft tests first draft

* implement tests first draft

* update fft/__init__.py

* expand tests fft

* expand tests ffftn

* expand tests and fix errors

* add Hermitian FFTs

* heat/fft/tests/test_fft.py

* raise IndexError, not ValueError, when axes don't match dimensions

* expand tests

* edit error message for better understanding

* replace == with allclose for 2D FFTs

* fix error

* remove redundant communication

* remove redundant tests

* fix bug in axes handling

* test hermitian FFT

* cast numpy fft2 to complex64

* expand tests

* edit error messages

* remove unnecessary axis check

* test inverse ffts as well

* skip comm-intensive tests on gpu

* introduce helper functions for real fft operations

* fix output shape wrt Nyquist frequency

* double precision tests

* debugging: introduce synchronization

* debugging

* debugging

* debugging

* specify default even size of last fft dim for inverse real ops

* add tests

* fix output shape calc when input is real

* expand tests

* expand tests

* expand tests

* add ihfft2, ihfft

* expand tests

* implement fftfreq, fftshift operations and tests

* fix local output dtype mismatch when local input tensor is empty

* remove print statements

* expand tests

* simplify dealing with multi-axis real FFT

* cannot be a list

* update documentation

* edit docs

---------

Co-authored-by: Fabian Hoppe <112093564+mrfh92@users.noreply.github.com>

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 *