Fix case & lint

llvm · Apr 15, 2024 · 890842c · 890842c
1 parent 05ab35d
commit 890842c
Show file tree

Hide file tree

Showing 4 changed files with 9 additions and 3 deletions.
diff --git a/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td b/include/torch-mlir/Dialect/Torch/IR/GeneratedTorchOps.td
@@ -15821,7 +15821,7 @@ def Torch_PrimsIotaOp : Torch_Op<"prims.iota", [
     Torch_BoolType:$requires_grad
   );
   let results = (outs
-    AnyTorchTensorType:$result
+    AnyTorchOptionalTensorType:$result
   );
   let hasCustomAssemblyFormat = 1;
   let extraClassDefinition = [{

diff --git a/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp b/lib/Dialect/Torch/Transforms/DecomposeComplexOps.cpp
@@ -4722,7 +4722,7 @@ class DecomposePrimsIotaOp : public OpRewritePattern<PrimsIotaOp> {
     return success();
   }
 };
-}  // namespace
+} // namespace
 
 namespace {
 // Decompose constant tensor full like ops.

diff --git a/projects/pt1/e2e_testing/xfail_sets.py b/projects/pt1/e2e_testing/xfail_sets.py
@@ -317,6 +317,9 @@
     "FloatImplicitModule_basic",
     "IntImplicitModule_basic",
 
+    # Unsupported: missing default value for argument 0 in schema for prims.iota.default
+    "PrimsIotaModule_basic",
+
     # Others
     "GridSamplerBasic1_basic",
     "GridSamplerBasic2_basic",
@@ -2438,6 +2441,9 @@
     # Failure - torch.aten.squeeze lower
     "BucketizeTensorOutInt32RightModule_basic", # unsupported by backend contract: tensor with unknown rank
 
+    # RuntimeError: unsupported input type: Device
+    "PrimsIotaModule_basic",
+
     # Failure - unknown
     "BucketizeTensorFloatModule_basic",
     "BucketizeTensorModule_basic",

diff --git a/projects/pt1/python/torch_mlir_e2e_test/test_suite/arange.py b/projects/pt1/python/torch_mlir_e2e_test/test_suite/arange.py
@@ -391,7 +391,7 @@ def __init__(self):
         None,
     ])
     def forward(self):
-        return torch.ops.prims.iota(77, start=0, step=1, dtype=torch.int64, device=torch.device('cpu'),
+        return torch.ops.prims.iota(77, start=0, step=1, dtype=torch.int64, device='cpu',
                                     requires_grad=False)
 
 @register_test_case(module_factory=lambda: PrimsIotaModule())