Copy2Map: formatted

spcl · Nov 26, 2023 · cfcdda8 · cfcdda8
1 parent 16a9969
commit cfcdda8
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 6 deletions.
diff --git a/dace/transformation/dataflow/copy_to_map.py b/dace/transformation/dataflow/copy_to_map.py
@@ -108,11 +108,17 @@ def apply(self, state: SDFGState, sdfg: SDFG):
 
         # Linearize and delinearize to get index expression for other side
         if copy_a:
-            a_index = [symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.src_subset.ranges[i][0]}') for i in range(len(copy_shape))]
+            a_index = [
+                symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.src_subset.ranges[i][0]}')
+                for i in range(len(copy_shape))
+            ]
             b_index = self.delinearize_linearize(bdesc, copy_shape, edge.data.get_dst_subset(edge, state))
         else:
             a_index = self.delinearize_linearize(adesc, copy_shape, edge.data.get_src_subset(edge, state))
-            b_index = [symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.dst_subset.ranges[i][0]}') for i in range(len(copy_shape))]
+            b_index = [
+                symbolic.pystr_to_symbolic(f'__i{i} + {edge.data.dst_subset.ranges[i][0]}')
+                for i in range(len(copy_shape))
+            ]
 
         a_subset = subsets.Range([(ind, ind, 1) for ind in a_index])
         b_subset = subsets.Range([(ind, ind, 1) for ind in b_index])

diff --git a/tests/transformations/copy_to_map_test.py b/tests/transformations/copy_to_map_test.py
@@ -7,6 +7,7 @@
 
 
 def _copy_to_map(storage: dace.StorageType):
+
     @dace
     def somecopy(a, b):
         b[:] = a
@@ -24,6 +25,7 @@ def somecopy(a, b):
 
 
 def _flatten_to_map(storage: dace.StorageType):
+
     @dace
     def somecopy(a, b):
         b[:] = a.flatten()
@@ -103,10 +105,11 @@ def test_preprocess():
 
 
 def test_squeezed_subset_src():
+
     @dace.program
-    def squeezed_subset(a: dace.float32[10], b: dace.float32[10,10]):
+    def squeezed_subset(a: dace.float32[10], b: dace.float32[10, 10]):
         for i in range(10):
-            b[i, i+1:10] = a[0:10-(i+1)]
+            b[i, i + 1:10] = a[0:10 - (i + 1)]
 
     sdfg = squeezed_subset.to_sdfg()
 
@@ -126,10 +129,11 @@ def squeezed_subset(a: dace.float32[10], b: dace.float32[10,10]):
 
 
 def test_squeezed_subset_src_swapped():
+
     @dace.program
-    def squeezed_subset(a: dace.float32[10], b: dace.float32[10,10]):
+    def squeezed_subset(a: dace.float32[10], b: dace.float32[10, 10]):
         for i in range(10):
-            b[i+1:10, i] = a[0:10-(i+1)]
+            b[i + 1:10, i] = a[0:10 - (i + 1)]
 
     sdfg = squeezed_subset.to_sdfg()