compiler: Fix decomposition dtype

ZoeLeibowitz · Oct 25, 2024 · 01af9bd · 01af9bd
1 parent 2ecf48a
commit 01af9bd
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/devito/mpi/distributed.py b/devito/mpi/distributed.py
@@ -259,8 +259,10 @@ def __init__(self, shape, dimensions, input_comm=None, topology=None):
             self._topology = tuple(1 for _ in range(len(shape)))
 
         # The domain decomposition
-        self._decomposition = [Decomposition(np.array_split(range(i), j), c)
-                               for i, j, c in zip(shape, self.topology, self.mycoords)]
+        self._decomposition = [
+            Decomposition(np.array_split(np.arange(i, dtype=d.dtype), j), c)
+            for i, j, c, d in zip(shape, self.topology, self.mycoords, dimensions)
+        ]
 
     @property
     def comm(self):