[mlir][Python] Add casting constructor to Type and Attribute.

[stellaraccident]

Summary

• This makes them consistent with custom types/attributes, whose constructors will do a type checked conversion. Of course, the base classes can represent everything so never error.
• More importantly, this makes it possible to subclass Type and Attribute out of tree in sensible ways.

[jdd]

LGTM! Thanks

[stellaraccident]

Rebase

mlir/lib/Bindings/Python/IRCore.cpp

py::class_<PyInsertionPoint>(m, "InsertionPoint")
                   py::arg("block"), "Inserts before the block terminator.")
       .def("insert", &PyInsertionPoint::insert, py::arg("operation"),
            "Inserts an operation.");
 
   //----------------------------------------------------------------------------
   // Mapping of PyAttribute.
   //----------------------------------------------------------------------------
   py::class_<PyAttribute>(m, "Attribute")
+      // Delegate to the PyAttribute copy constructor, which will also lifetime
+      // extend the backing context which owns the MlirAttribute.
+      .def(py::init<PyAttribute &>(), py::arg("cast_from_type"),
+           "Casts the passed attribute to the generic Attribute")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR,
                              &PyAttribute::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyAttribute::createFromCapsule)
       .def_static(
           "parse",
           [](std::string attrSpec, DefaultingPyMlirContext context) {
             MlirAttribute type = mlirAttributeParseGet(
                 context->get(), toMlirStringRef(attrSpec));
py::class_<PyNamedAttribute>(m, "NamedAttribute")
           },
           py::keep_alive<0, 1>(),
           "The underlying generic attribute of the NamedAttribute binding");
 
   //----------------------------------------------------------------------------
   // Mapping of PyType.
   //----------------------------------------------------------------------------
   py::class_<PyType>(m, "Type")
+      // Delegate to the PyType copy constructor, which will also lifetime
+      // extend the backing context which owns the MlirType.
+      .def(py::init<PyType &>(), py::arg("cast_from_type"),
+           "Casts the passed type to the generic Type")
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyType::getCapsule)
       .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyType::createFromCapsule)
       .def_static(
           "parse",
           [](std::string typeSpec, DefaultingPyMlirContext context) {
             MlirType type =
                 mlirTypeParseGet(context->get(), toMlirStringRef(typeSpec));
             // TODO: Rework error reporting once diagnostic engine is exposed

mlir/test/Bindings/Python/ir_attributes.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testParsePrint
+@run
 def testParsePrint():
   with Context() as ctx:
     t = Attribute.parse('"hello"')
   assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: "hello"
   print(str(t))
   # CHECK: Attribute("hello")
   print(repr(t))
 
-run(testParsePrint)
-
 
 # CHECK-LABEL: TEST: testParseError
 # TODO: Hook the diagnostic manager to capture a more meaningful error
 # message.
+@run
 def testParseError():
   with Context():
     try:
       t = Attribute.parse("BAD_ATTR_DOES_NOT_EXIST")
     except ValueError as e:
       # CHECK: Unable to parse attribute: 'BAD_ATTR_DOES_NOT_EXIST'
       print("testParseError:", e)
     else:
       print("Exception not produced")
 
-run(testParseError)
-
 
 # CHECK-LABEL: TEST: testAttrEq
+@run
 def testAttrEq():
   with Context():
     a1 = Attribute.parse('"attr1"')
     a2 = Attribute.parse('"attr2"')
     a3 = Attribute.parse('"attr1"')
     # CHECK: a1 == a1: True
     print("a1 == a1:", a1 == a1)
     # CHECK: a1 == a2: False
     print("a1 == a2:", a1 == a2)
     # CHECK: a1 == a3: True
     print("a1 == a3:", a1 == a3)
     # CHECK: a1 == None: False
     print("a1 == None:", a1 == None)
 
-run(testAttrEq)
+
+# CHECK-LABEL: TEST: testAttrCast
+@run
+def testAttrCast():
+  with Context():
+    a1 = Attribute.parse('"attr1"')
+    a2 = Attribute(a1)
+    # CHECK: a1 == a2: True
+    print("a1 == a2:", a1 == a2)
 
 
 # CHECK-LABEL: TEST: testAttrEqDoesNotRaise
+@run
 def testAttrEqDoesNotRaise():
   with Context():
     a1 = Attribute.parse('"attr1"')
     not_an_attr = "foo"
     # CHECK: False
     print(a1 == not_an_attr)
     # CHECK: False
     print(a1 == None)
     # CHECK: True
     print(a1 != None)
 
-run(testAttrEqDoesNotRaise)
-
 
 # CHECK-LABEL: TEST: testAttrCapsule
+@run
 def testAttrCapsule():
   with Context() as ctx:
     a1 = Attribute.parse('"attr1"')
   # CHECK: mlir.ir.Attribute._CAPIPtr
   attr_capsule = a1._CAPIPtr
   print(attr_capsule)
   a2 = Attribute._CAPICreate(attr_capsule)
   assert a2 == a1
   assert a2.context is ctx
 
-run(testAttrCapsule)
-
 
 # CHECK-LABEL: TEST: testStandardAttrCasts
+@run
 def testStandardAttrCasts():
   with Context():
     a1 = Attribute.parse('"attr1"')
     astr = StringAttr(a1)
     aself = StringAttr(astr)
     # CHECK: Attribute("attr1")
     print(repr(astr))
     try:
       tillegal = StringAttr(Attribute.parse("1.0"))
     except ValueError as e:
       # CHECK: ValueError: Cannot cast attribute to StringAttr (from Attribute(1.000000e+00 : f64))
       print("ValueError:", e)
     else:
       print("Exception not produced")
 
-run(testStandardAttrCasts)
-
 
 # CHECK-LABEL: TEST: testAffineMapAttr
+@run
 def testAffineMapAttr():
   with Context() as ctx:
     d0 = AffineDimExpr.get(0)
     d1 = AffineDimExpr.get(1)
     c2 = AffineConstantExpr.get(2)
     map0 = AffineMap.get(2, 3, [])
 
     # CHECK: affine_map<(d0, d1)[s0, s1, s2] -> ()>
     attr_built = AffineMapAttr.get(map0)
     print(str(attr_built))
 
     attr_parsed = Attribute.parse(str(attr_built))
     assert attr_built == attr_parsed
 
-run(testAffineMapAttr)
-
 
 # CHECK-LABEL: TEST: testFloatAttr
+@run
 def testFloatAttr():
   with Context(), Location.unknown():
     fattr = FloatAttr(Attribute.parse("42.0 : f32"))
     # CHECK: fattr value: 42.0
     print("fattr value:", fattr.value)
 
     # Test factory methods.
     # CHECK: default_get: 4.200000e+01 : f32
     print("default_get:", FloatAttr.get(
         F32Type.get(), 42.0))
     # CHECK: f32_get: 4.200000e+01 : f32
     print("f32_get:", FloatAttr.get_f32(42.0))
     # CHECK: f64_get: 4.200000e+01 : f64
     print("f64_get:", FloatAttr.get_f64(42.0))
     try:
       fattr_invalid = FloatAttr.get(
           IntegerType.get_signless(32), 42)
     except ValueError as e:
       # CHECK: invalid 'Type(i32)' and expected floating point type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testFloatAttr)
-
 
 # CHECK-LABEL: TEST: testIntegerAttr
+@run
 def testIntegerAttr():
   with Context() as ctx:
     iattr = IntegerAttr(Attribute.parse("42"))
     # CHECK: iattr value: 42
     print("iattr value:", iattr.value)
     # CHECK: iattr type: i64
     print("iattr type:", iattr.type)
 
     # Test factory methods.
     # CHECK: default_get: 42 : i32
     print("default_get:", IntegerAttr.get(
         IntegerType.get_signless(32), 42))
 
-run(testIntegerAttr)
-
 
 # CHECK-LABEL: TEST: testBoolAttr
+@run
 def testBoolAttr():
   with Context() as ctx:
     battr = BoolAttr(Attribute.parse("true"))
     # CHECK: iattr value: True
     print("iattr value:", battr.value)
 
     # Test factory methods.
     # CHECK: default_get: true
     print("default_get:", BoolAttr.get(True))
 
-run(testBoolAttr)
-
 
 # CHECK-LABEL: TEST: testFlatSymbolRefAttr
+@run
 def testFlatSymbolRefAttr():
   with Context() as ctx:
     sattr = FlatSymbolRefAttr(Attribute.parse('@symbol'))
     # CHECK: symattr value: symbol
     print("symattr value:", sattr.value)
 
     # Test factory methods.
     # CHECK: default_get: @foobar
     print("default_get:", FlatSymbolRefAttr.get("foobar"))
 
-run(testFlatSymbolRefAttr)
-
 
 # CHECK-LABEL: TEST: testStringAttr
+@run
 def testStringAttr():
   with Context() as ctx:
     sattr = StringAttr(Attribute.parse('"stringattr"'))
     # CHECK: sattr value: stringattr
     print("sattr value:", sattr.value)
 
     # Test factory methods.
     # CHECK: default_get: "foobar"
     print("default_get:", StringAttr.get("foobar"))
     # CHECK: typed_get: "12345" : i32
     print("typed_get:", StringAttr.get_typed(
         IntegerType.get_signless(32), "12345"))
 
-run(testStringAttr)
-
 
 # CHECK-LABEL: TEST: testNamedAttr
+@run
 def testNamedAttr():
   with Context():
     a = Attribute.parse('"stringattr"')
     named = a.get_named("foobar")  # Note: under the small object threshold
     # CHECK: attr: "stringattr"
     print("attr:", named.attr)
     # CHECK: name: foobar
     print("name:", named.name)
     # CHECK: named: NamedAttribute(foobar="stringattr")
     print("named:", named)
 
-run(testNamedAttr)
-
 
 # CHECK-LABEL: TEST: testDenseIntAttr
+@run
 def testDenseIntAttr():
   with Context():
     raw = Attribute.parse("dense<[[0,1,2],[3,4,5]]> : vector<2x3xi32>")
     # CHECK: attr: dense<[{{\[}}0, 1, 2], [3, 4, 5]]>
     print("attr:", raw)
 
     a = DenseIntElementsAttr(raw)
     assert len(a) == 6
with Context():
     for value in a:
       print(value, end=" ")
     print()
 
     # CHECK: i1
     print(ShapedType(a.type).element_type)
 
 
-run(testDenseIntAttr)
-
-
 # CHECK-LABEL: TEST: testDenseFPAttr
+@run
 def testDenseFPAttr():
   with Context():
     raw = Attribute.parse("dense<[0.0, 1.0, 2.0, 3.0]> : vector<4xf32>")
     # CHECK: attr: dense<[0.000000e+00, 1.000000e+00, 2.000000e+00, 3.000000e+00]>
 
     print("attr:", raw)
 
     a = DenseFPElementsAttr(raw)
     assert len(a) == 4
 
     # CHECK: 0.0 1.0 2.0 3.0
     for value in a:
       print(value, end=" ")
     print()
 
     # CHECK: f32
     print(ShapedType(a.type).element_type)
 
 
-run(testDenseFPAttr)
-
-
 # CHECK-LABEL: TEST: testDictAttr
+@run
 def testDictAttr():
   with Context():
     dict_attr = {
       'stringattr':  StringAttr.get('string'),
       'integerattr' : IntegerAttr.get(
         IntegerType.get_signless(32), 42)
     }
 
with Context():
     try:
       _ = a[42]
     except IndexError:
       pass
     else:
       assert False, "expected IndexError on accessing an out-of-bounds attribute"
 
 
-
-run(testDictAttr)
-
 # CHECK-LABEL: TEST: testTypeAttr
+@run
 def testTypeAttr():
   with Context():
     raw = Attribute.parse("vector<4xf32>")
     # CHECK: attr: vector<4xf32>
     print("attr:", raw)
     type_attr = TypeAttr(raw)
     # CHECK: f32
     print(ShapedType(type_attr.value).element_type)
 
 
-run(testTypeAttr)
-
-
 # CHECK-LABEL: TEST: testArrayAttr
+@run
 def testArrayAttr():
   with Context():
     raw = Attribute.parse("[42, true, vector<4xf32>]")
   # CHECK: attr: [42, true, vector<4xf32>]
   print("raw attr:", raw)
   # CHECK: - 42
   # CHECK: - true
   # CHECK: - vector<4xf32>
with Context():
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute (None?) when attempting to create an ArrayAttribute
       print("Error: ", e)
     try:
       ArrayAttr.get([42])
     except RuntimeError as e:
       # CHECK: Error: Invalid attribute when attempting to create an ArrayAttribute
       print("Error: ", e)
-run(testArrayAttr)
 

mlir/test/Bindings/Python/ir_types.py

 # RUN: %PYTHON %s | FileCheck %s
 
 import gc
 from mlir.ir import *
 
 def run(f):
   print("\nTEST:", f.__name__)
   f()
   gc.collect()
   assert Context._get_live_count() == 0
+  return f
 
 
 # CHECK-LABEL: TEST: testParsePrint
+@run
 def testParsePrint():
   ctx = Context()
   t = Type.parse("i32", ctx)
   assert t.context is ctx
   ctx = None
   gc.collect()
   # CHECK: i32
   print(str(t))
   # CHECK: Type(i32)
   print(repr(t))
 
-run(testParsePrint)
-
 
 # CHECK-LABEL: TEST: testParseError
 # TODO: Hook the diagnostic manager to capture a more meaningful error
 # message.
+@run
 def testParseError():
   ctx = Context()
   try:
     t = Type.parse("BAD_TYPE_DOES_NOT_EXIST", ctx)
   except ValueError as e:
     # CHECK: Unable to parse type: 'BAD_TYPE_DOES_NOT_EXIST'
     print("testParseError:", e)
   else:
     print("Exception not produced")
 
-run(testParseError)
-
 
 # CHECK-LABEL: TEST: testTypeEq
+@run
 def testTypeEq():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type.parse("f32", ctx)
   t3 = Type.parse("i32", ctx)
   # CHECK: t1 == t1: True
   print("t1 == t1:", t1 == t1)
   # CHECK: t1 == t2: False
   print("t1 == t2:", t1 == t2)
   # CHECK: t1 == t3: True
   print("t1 == t3:", t1 == t3)
   # CHECK: t1 == None: False
   print("t1 == None:", t1 == None)
 
-run(testTypeEq)
+
+# CHECK-LABEL: TEST: testTypeCast
+@run
+def testTypeCast():
+  ctx = Context()
+  t1 = Type.parse("i32", ctx)
+  t2 = Type(t1)
+  # CHECK: t1 == t2: True
+  print("t1 == t2:", t1 == t2)
 
 
 # CHECK-LABEL: TEST: testTypeIsInstance
+@run
 def testTypeIsInstance():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   t2 = Type.parse("f32", ctx)
   # CHECK: True
   print(IntegerType.isinstance(t1))
   # CHECK: False
   print(F32Type.isinstance(t1))
   # CHECK: True
   print(F32Type.isinstance(t2))
 
-run(testTypeIsInstance)
-
 
 # CHECK-LABEL: TEST: testTypeEqDoesNotRaise
+@run
 def testTypeEqDoesNotRaise():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   not_a_type = "foo"
   # CHECK: False
   print(t1 == not_a_type)
   # CHECK: False
   print(t1 == None)
   # CHECK: True
   print(t1 != None)
 
-run(testTypeEqDoesNotRaise)
-
 
 # CHECK-LABEL: TEST: testTypeCapsule
+@run
 def testTypeCapsule():
   with Context() as ctx:
     t1 = Type.parse("i32", ctx)
   # CHECK: mlir.ir.Type._CAPIPtr
   type_capsule = t1._CAPIPtr
   print(type_capsule)
   t2 = Type._CAPICreate(type_capsule)
   assert t2 == t1
   assert t2.context is ctx
 
-run(testTypeCapsule)
-
 
 # CHECK-LABEL: TEST: testStandardTypeCasts
+@run
 def testStandardTypeCasts():
   ctx = Context()
   t1 = Type.parse("i32", ctx)
   tint = IntegerType(t1)
   tself = IntegerType(tint)
   # CHECK: Type(i32)
   print(repr(tint))
   try:
     tillegal = IntegerType(Type.parse("f32", ctx))
   except ValueError as e:
     # CHECK: ValueError: Cannot cast type to IntegerType (from Type(f32))
     print("ValueError:", e)
   else:
     print("Exception not produced")
 
-run(testStandardTypeCasts)
-
 
 # CHECK-LABEL: TEST: testIntegerType
+@run
 def testIntegerType():
   with Context() as ctx:
     i32 = IntegerType(Type.parse("i32"))
     # CHECK: i32 width: 32
     print("i32 width:", i32.width)
     # CHECK: i32 signless: True
     print("i32 signless:", i32.is_signless)
     # CHECK: i32 signed: False
with Context() as ctx:
 
     # CHECK: signless: i16
     print("signless:", IntegerType.get_signless(16))
     # CHECK: signed: si8
     print("signed:", IntegerType.get_signed(8))
     # CHECK: unsigned: ui64
     print("unsigned:", IntegerType.get_unsigned(64))
 
-run(testIntegerType)
-
 # CHECK-LABEL: TEST: testIndexType
+@run
 def testIndexType():
   with Context() as ctx:
     # CHECK: index type: index
     print("index type:", IndexType.get())
 
-run(testIndexType)
 
 # CHECK-LABEL: TEST: testFloatType
+@run
 def testFloatType():
   with Context():
     # CHECK: float: bf16
     print("float:", BF16Type.get())
     # CHECK: float: f16
     print("float:", F16Type.get())
     # CHECK: float: f32
     print("float:", F32Type.get())
     # CHECK: float: f64
     print("float:", F64Type.get())
 
-run(testFloatType)
 
 # CHECK-LABEL: TEST: testNoneType
+@run
 def testNoneType():
   with Context():
     # CHECK: none type: none
     print("none type:", NoneType.get())
 
-run(testNoneType)
 
 # CHECK-LABEL: TEST: testComplexType
+@run
 def testComplexType():
   with Context() as ctx:
     complex_i32 = ComplexType(Type.parse("complex<i32>"))
     # CHECK: complex type element: i32
     print("complex type element:", complex_i32.element_type)
 
     f32 = F32Type.get()
     # CHECK: complex type: complex<f32>
     print("complex type:", ComplexType.get(f32))
 
     index = IndexType.get()
     try:
       complex_invalid = ComplexType.get(index)
     except ValueError as e:
       # CHECK: invalid 'Type(index)' and expected floating point or integer type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testComplexType)
-
 
 # CHECK-LABEL: TEST: testConcreteShapedType
 # Shaped type is not a kind of builtin types, it is the base class for vectors,
 # memrefs and tensors, so this test case uses an instance of vector to test the
 # shaped type. The class hierarchy is preserved on the python side.
+@run
 def testConcreteShapedType():
   with Context() as ctx:
     vector = VectorType(Type.parse("vector<2x3xf32>"))
     # CHECK: element type: f32
     print("element type:", vector.element_type)
     # CHECK: whether the given shaped type is ranked: True
     print("whether the given shaped type is ranked:", vector.has_rank)
     # CHECK: rank: 2
     print("rank:", vector.rank)
     # CHECK: whether the shaped type has a static shape: True
     print("whether the shaped type has a static shape:", vector.has_static_shape)
     # CHECK: whether the dim-th dimension is dynamic: False
     print("whether the dim-th dimension is dynamic:", vector.is_dynamic_dim(0))
     # CHECK: dim size: 3
     print("dim size:", vector.get_dim_size(1))
     # CHECK: is_dynamic_size: False
     print("is_dynamic_size:", vector.is_dynamic_size(3))
     # CHECK: is_dynamic_stride_or_offset: False
     print("is_dynamic_stride_or_offset:", vector.is_dynamic_stride_or_offset(1))
     # CHECK: isinstance(ShapedType): True
     print("isinstance(ShapedType):", isinstance(vector, ShapedType))
 
-run(testConcreteShapedType)
 
 # CHECK-LABEL: TEST: testAbstractShapedType
 # Tests that ShapedType operates as an abstract base class of a concrete
 # shaped type (using vector as an example).
+@run
 def testAbstractShapedType():
   ctx = Context()
   vector = ShapedType(Type.parse("vector<2x3xf32>", ctx))
   # CHECK: element type: f32
   print("element type:", vector.element_type)
 
-run(testAbstractShapedType)
 
 # CHECK-LABEL: TEST: testVectorType
+@run
 def testVectorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     # CHECK: vector type: vector<2x3xf32>
     print("vector type:", VectorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       vector_invalid = VectorType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point or integer type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testVectorType)
 
 # CHECK-LABEL: TEST: testRankedTensorType
+@run
 def testRankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
     # CHECK: ranked tensor type: tensor<2x3xf32>
     print("ranked tensor type:",
           RankedTensorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       tensor_invalid = RankedTensorType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testRankedTensorType)
 
 # CHECK-LABEL: TEST: testUnrankedTensorType
+@run
 def testUnrankedTensorType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
     unranked_tensor = UnrankedTensorType.get(f32)
     # CHECK: unranked tensor type: tensor<*xf32>
     print("unranked tensor type:", unranked_tensor)
     try:
try:
       tensor_invalid = UnrankedTensorType.get(none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testUnrankedTensorType)
 
 # CHECK-LABEL: TEST: testMemRefType
+@run
 def testMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     shape = [2, 3]
     loc = Location.unknown()
     memref = MemRefType.get(shape, f32, memory_space=Attribute.parse("2"))
     # CHECK: memref type: memref<2x3xf32, 2>
     print("memref type:", memref)
try:
       memref_invalid = MemRefType.get(shape, none)
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testMemRefType)
 
 # CHECK-LABEL: TEST: testUnrankedMemRefType
+@run
 def testUnrankedMemRefType():
   with Context(), Location.unknown():
     f32 = F32Type.get()
     loc = Location.unknown()
     unranked_memref = UnrankedMemRefType.get(f32, Attribute.parse("2"))
     # CHECK: unranked memref type: memref<*xf32, 2>
     print("unranked memref type:", unranked_memref)
     try:
try:
       memref_invalid = UnrankedMemRefType.get(none, Attribute.parse("2"))
     except ValueError as e:
       # CHECK: invalid 'Type(none)' and expected floating point, integer, vector
       # CHECK: or complex type.
       print(e)
     else:
       print("Exception not produced")
 
-run(testUnrankedMemRefType)
 
 # CHECK-LABEL: TEST: testTupleType
+@run
 def testTupleType():
   with Context() as ctx:
     i32 = IntegerType(Type.parse("i32"))
     f32 = F32Type.get()
     vector = VectorType(Type.parse("vector<2x3xf32>"))
     l = [i32, f32, vector]
     tuple_type = TupleType.get_tuple(l)
     # CHECK: tuple type: tuple<i32, f32, vector<2x3xf32>>
     print("tuple type:", tuple_type)
     # CHECK: number of types: 3
     print("number of types:", tuple_type.num_types)
     # CHECK: pos-th type in the tuple type: f32
     print("pos-th type in the tuple type:", tuple_type.get_type(1))
 
-run(testTupleType)
-
 
 # CHECK-LABEL: TEST: testFunctionType
+@run
 def testFunctionType():
   with Context() as ctx:
     input_types = [IntegerType.get_signless(32),
                   IntegerType.get_signless(16)]
     result_types = [IndexType.get()]
     func = FunctionType.get(input_types, result_types)
     # CHECK: INPUTS: [Type(i32), Type(i16)]
     print("INPUTS:", func.inputs)
     # CHECK: RESULTS: [Type(index)]
     print("RESULTS:", func.results)
-
-
-run(testFunctionType)