[Mlir-commits] [mlir] 91865cc - [mlir][taco] Accept an integer list for the ordering when defining a tensor format.

[Bixia Zheng]

Author: Bixia Zheng
Date: 2022-01-28T10:33:25-08:00
New Revision: 91865cc027612e208ed30c5f6b107d0c87c59193

URL: https://github.com/llvm/llvm-project/commit/91865cc027612e208ed30c5f6b107d0c87c59193
DIFF: https://github.com/llvm/llvm-project/commit/91865cc027612e208ed30c5f6b107d0c87c59193.diff

LOG: [mlir][taco] Accept an integer list for the ordering when defining a tensor format.

The unit tests for PyTACO hasn't been upstreamed yet. A unit test for this
change will be added when we upstream all the unit tests for PyTACO.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D118417

Added: 
    

Modified: 
    mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py

Removed: 
    


################################################################################
diff  --git a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
index f64d34037eabd..657185cd0350a 100644
--- a/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
+++ b/mlir/test/Integration/Dialect/SparseTensor/taco/tools/mlir_pytaco.py
@@ -323,6 +323,10 @@ def __post_init__(self) -> None:
 
     if self.ordering is None:
       self.ordering = ModeOrdering(list(range(self.rank())))
+    if isinstance(self.ordering, list):
+      if not _all_instance_of(self.ordering, int):
+        raise ValueError(f"Expected a list of integer: {self.ordering}")
+      self.ordering = ModeOrdering(self.ordering)
     if not isinstance(self.ordering, ModeOrdering):
       raise ValueError(f"Expected ModeOrdering: {self.ordering}")