Fix partitioner bug

FPSG-UIUC · Sep 24, 2024 · ea99815 · ea99815
1 parent b5d5c53
commit ea99815
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/teaal/trans/partitioner.py b/teaal/trans/partitioner.py
@@ -110,9 +110,12 @@ def unpartition(self, tensor: Tensor) -> Statement:
  swizzled_ranks, part_ir.get_all_parts(), False)
  for part in valid_parts:
  trans.append((part, tensor.get_ranks()))
+
+ # If this is a partition, apply all partitioning
+ # If this is a flatten, just flatten
  tensor.update_ranks(
  part_ir.partition_ranks(
- tensor.get_ranks(), {part}, True, False))
+ tensor.get_ranks(), {part}, len(part) == 1, False))
 
  new_ranks = tensor.get_ranks()