Synchronize the collection of symbolic torch ops

llvm · Apr 26, 2024 · dfd311c · dfd311c
1 parent 634a796
commit dfd311c
Showing 1 changed file with 7 additions and 11 deletions.
diff --git a/python/torch_mlir/extras/fx_importer.py b/python/torch_mlir/extras/fx_importer.py
@@ -236,26 +236,18 @@
 # set and just check key existence in SYMBOLIC_OP_TO_TORCH_OP
 
 if _IS_TORCH_2_1_OR_EARLIER:
-    SYMBOLIC_TORCH_OPS = {
-        torch.ops.aten.sym_size,
-        torch.ops.aten.sym_stride,
-        torch.ops.aten.sym_numel,
-    }
-
     SYMBOLIC_OP_TO_TORCH_OP = {
         (torch.ops.aten.sym_size, 1): torch.ops.aten.size.default,
         (torch.ops.aten.sym_size, 2): torch.ops.aten.size.int,
         (torch.ops.aten.sym_stride, 1): torch.ops.aten.stride.default,
         (torch.ops.aten.sym_stride, 2): torch.ops.aten.stride.int,
         (torch.ops.aten.sym_numel, 1): torch.ops.aten.numel.default,
     }
-else:
+
     SYMBOLIC_TORCH_OPS = {
-        torch.ops.aten.sym_size.int,
-        torch.ops.aten.sym_stride.int,
-        torch.ops.aten.sym_numel.default,
+        x[0] for x in SYMBOLIC_OP_TO_TORCH_OP
     }
-
+else:
     SYMBOLIC_OP_TO_TORCH_OP = {
         torch.ops.aten.sym_size.default: torch.ops.aten.size.default,
         torch.ops.aten.sym_size.int: torch.ops.aten.size.int,
@@ -264,6 +256,10 @@
         torch.ops.aten.sym_numel.default: torch.ops.aten.numel.default,
     }
 
+    SYMBOLIC_TORCH_OPS = {
+        x for x in SYMBOLIC_OP_TO_TORCH_OP
+    }
+
 
 @dataclass(frozen=True)
 class SparsityMeta: