[Op] Enforce int64 output shape in CallTIR (tlc-pack#322)

junrushao · Jan 25, 2023 · e4eb457 · e4eb457
1 parent d079eee
commit e4eb457
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/python/tvm/relax/op/base.py b/python/tvm/relax/op/base.py
@@ -68,10 +68,11 @@ def _create_shape(shape: List[Union[int, PrimExpr]]) -> ShapeExpr:
         for x in shape:
             if isinstance(x, int):
                 shape_array.append(tvm.tir.IntImm("int64", x))
+            elif isinstance(x, tvm.tir.IntImm):
+                shape_array.append(x if x.dtype == "int64" else tvm.tir.IntImm("int64", x.value))
             elif isinstance(x, PrimExpr):
-                # TODO: enforce all shapes are i64
-                # if x.dtype != "int64":
-                #     raise TypeError("Expect int64 dtype for shape")
+                if x.dtype != "int64":
+                    raise TypeError("Expect int64 dtype for shape")
                 shape_array.append(x)
             else:
                 raise TypeError("Expect int or PrimExpr for shape")