[MLIR][python bindings] Expose TypeIDs in python

[makslevental]

Summary

This diff adds python bindings for MlirTypeID. It paves the way for returning accurately typed Types from python APIs (see D150927) and then further along building type "conscious" Value APIs (see D150413).

[makslevental]

remove CAPI tests

[makslevental]

fix cruft

[makslevental]

update commit message

[makslevental]

add type_caster for MlirTypeID

[makslevental]

done

[rkayaith]

mlir/lib/Bindings/Python/IRCore.cpp

3303

In other places I see "cheap" functions like this are usually exposed as properties. There's also the class-level variant def_property_readonly_static.

mlir/lib/Bindings/Python/IRModule.h

897

Can you add this to mlir_type_subclass as well, and implement it for some non-builtin types as an example.

900

Instead of returning None, could we just skip defining the function? That would make it obvious when a class doesn't have it implemented.

mlir/test/python/ir/builtin_types.py

571–573

maybe just define these as a single list of tuples?

626

I would expect get_typeid() here to work the same as Type.get_typeid(), and return the concrete typeid (the id for vector in this case).
Maybe rename the static methods (static_typeid?) to disambiguate

[makslevental]

incorporate comments

mlir/lib/Bindings/Python/IRModule.h

900

I don't understand why (some pybind impl detail) but skipping defining on one such concrete leaves a property object descriptor there; i.e.,

if (DerivedTy::getTypeIdFunction)
  cls.def_property_readonly_static(
      "typeid", [](py::object & /*class*/) -> MlirTypeID {
        return DerivedTy::getTypeIdFunction();
      });

results in <property object at 0x7efe2dcebc40> being printed for ShapeType.typeid. Could do

if (DerivedTy::getTypeIdFunction)
  cls.def_property_readonly_static(
      "typeid", [](py::object & /*class*/) -> MlirTypeID {
        return DerivedTy::getTypeIdFunction();
      });
else
  cls.attr("typeid") = py::none();

but that just looks the same from the user side and is baroque (imho). Changed to throw a ValueError.

[makslevental]

add static_typeid to mlir_type_subclass

[jpienaar]

mlir/include/mlir-c/BuiltinTypes.h

63

OOC why not just

MlirTypeID mlirTypeGetTypeID(MlirType type);

?

[makslevental]

remove unnecessary def_property_readonly_static

[makslevental]

mlir/include/mlir-c/BuiltinTypes.h

63

You mean in general why did I add all of these () -> TypeID functions? Because in the TypeCaster diff you need to eg build a map from TypeID -> TypeCaster without having a type on hand.

[rkayaith]

mlir/include/mlir-c/BuiltinTypes.h

63

if only we could generate all the C api boilerplate :)

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

3309–3310

Let's use C++ API: throw py::value_error(...). Here and below.

[makslevental]

mlir/lib/Bindings/Python/IRCore.cpp

3309–3310

Just clarifying where "below" refers to: there are ~50 instances of throw SetPyError in this directory - do you want me to replace all of them? If so, no problem, but that would be better in a separate diff? Otherwise ctrl+f in this UI, I see this one and the one in PyConcreteType::castFrom.

[ftynse]

mlir/lib/Bindings/Python/IRCore.cpp

3309–3310

I meant only the ones introduced in this diff. Though changing all the others in a separate diff is also much appreciated.

[makslevental]

use C++ API throw py::value_error

[rkayaith]

mlir/include/mlir-c/BuiltinTypes.h

76

nit, same for a bunch of others as well

mlir/test/python/ir/builtin_types.py

580–582

the keys of the dict here are the python *Type classes (IntegerType etc.), not PyTypeIDs like the comment says

mlir/include/mlir-c/Bindings/Python/Interop.h

 #define MLIR_PYTHON_CAPSULE_MODULE                                             \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Module._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_OPERATION                                          \
   MAKE_MLIR_PYTHON_QUALNAME("ir.Operation._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_TYPE MAKE_MLIR_PYTHON_QUALNAME("ir.Type._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_PASS_MANAGER                                       \
   MAKE_MLIR_PYTHON_QUALNAME("passmanager.PassManager._CAPIPtr")
 #define MLIR_PYTHON_CAPSULE_VALUE MAKE_MLIR_PYTHON_QUALNAME("ir.Value._CAPIPtr")
+#define MLIR_PYTHON_CAPSULE_TYPEID                                             \
+  MAKE_MLIR_PYTHON_QUALNAME("ir.TypeID._CAPIPtr")
 
 /** Attribute on MLIR Python objects that expose their C-API pointer.
  * This will be a type-specific capsule created as per one of the helpers
  * below.
  *
  * Ownership is not transferred by acquiring a capsule in this way: the
  * validity of the pointer wrapped by the capsule will be bounded by the
  * lifetime of the Python object that produced it. Only the name and pointer

  * a null type is returned (as checked via mlirOperationIsNull). In such a
  * case, the Python APIs will have already set an error. */
 static inline MlirOperation mlirPythonCapsuleToOperation(PyObject *capsule) {
   void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_OPERATION);
   MlirOperation op = {ptr};
   return op;
 }
 
+/** Creates a capsule object encapsulating the raw C-API MlirTypeID.
+ * The returned capsule does not extend or affect ownership of any Python
+ * objects that reference the type in any way.
+ */
+static inline PyObject *mlirPythonTypeIDToCapsule(MlirTypeID typeID) {
+  return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(typeID),
+                       MLIR_PYTHON_CAPSULE_TYPEID, NULL);
+}
+
+/** Extracts an MlirTypeID from a capsule as produced from
+ * mlirPythonTypeIDToCapsule. If the capsule is not of the right type, then
+ * a null type is returned (as checked via mlirTypeIDIsNull). In such a
+ * case, the Python APIs will have already set an error. */
+static inline MlirTypeID mlirPythonCapsuleToTypeID(PyObject *capsule) {
+  void *ptr = PyCapsule_GetPointer(capsule, MLIR_PYTHON_CAPSULE_TYPEID);
+  MlirTypeID typeID = {ptr};
+  return typeID;
+}
+
 /** Creates a capsule object encapsulating the raw C-API MlirType.
  * The returned capsule does not extend or affect ownership of any Python
  * objects that reference the type in any way.
  */
 static inline PyObject *mlirPythonTypeToCapsule(MlirType type) {
   return PyCapsule_New(MLIR_PYTHON_GET_WRAPPED_POINTER(type),
                        MLIR_PYTHON_CAPSULE_TYPE, NULL);
 }

mlir/include/mlir-c/BuiltinTypes.h

 #ifdef __cplusplus
 extern "C" {
 #endif
 
 //===----------------------------------------------------------------------===//
 // Integer types.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Integer type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirIntegerTypeGetTypeID(void);
+
 /// Checks whether the given type is an integer type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAInteger(MlirType type);
 
 /// Creates a signless integer type of the given bitwidth in the context. The
 /// type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirIntegerTypeGet(MlirContext ctx,
                                                unsigned bitwidth);
 

 
 /// Checks whether the given integer type is unsigned.
 MLIR_CAPI_EXPORTED bool mlirIntegerTypeIsUnsigned(MlirType type);
 
 //===----------------------------------------------------------------------===//
 // Index type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Index type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirIndexTypeGetTypeID(void);
+
 /// Checks whether the given type is an index type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAIndex(MlirType type);
 
 /// Creates an index type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirIndexTypeGet(MlirContext ctx);
 
 //===----------------------------------------------------------------------===//
 // Floating-point types.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Float8E5M2 type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M2TypeGetTypeID(void);
+
 /// Checks whether the given type is an f8E5M2 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M2(MlirType type);
 
 /// Creates an f8E5M2 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M2TypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float8E4M3FN type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E4M3FNTypeGetTypeID(void);
+
 /// Checks whether the given type is an f8E4M3FN type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3FN(MlirType type);
 
 /// Creates an f8E4M3FN type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float8E5M2FNUZ type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E5M2FNUZTypeGetTypeID(void);
+
 /// Checks whether the given type is an f8E5M2FNUZ type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E5M2FNUZ(MlirType type);
 
 /// Creates an f8E5M2FNUZ type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E5M2FNUZTypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float8E4M3FNUZ type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E4M3FNUZTypeGetTypeID(void);
+
 /// Checks whether the given type is an f8E4M3FNUZ type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3FNUZ(MlirType type);
 
 /// Creates an f8E4M3FNUZ type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3FNUZTypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float8E4M3B11FNUZ type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat8E4M3B11FNUZTypeGetTypeID(void);
+
 /// Checks whether the given type is an f8E4M3B11FNUZ type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFloat8E4M3B11FNUZ(MlirType type);
 
 /// Creates an f8E4M3B11FNUZ type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirFloat8E4M3B11FNUZTypeGet(MlirContext ctx);
 
+/// Returns the typeID of an BFloat16 type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirBFloat16TypeGetTypeID(void);
+
 /// Checks whether the given type is a bf16 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsABF16(MlirType type);
 
 /// Creates a bf16 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirBF16TypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float16 type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat16TypeGetTypeID(void);
+
 /// Checks whether the given type is an f16 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAF16(MlirType type);
 
 /// Creates an f16 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirF16TypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float32 type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat32TypeGetTypeID(void);
+
 /// Checks whether the given type is an f32 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAF32(MlirType type);
 
 /// Creates an f32 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirF32TypeGet(MlirContext ctx);
 
+/// Returns the typeID of an Float64 type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFloat64TypeGetTypeID(void);
+
 /// Checks whether the given type is an f64 type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAF64(MlirType type);
 
 /// Creates a f64 type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirF64TypeGet(MlirContext ctx);
 
 //===----------------------------------------------------------------------===//
 // None type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an None type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirNoneTypeGetTypeID(void);
+
 /// Checks whether the given type is a None type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsANone(MlirType type);
 
 /// Creates a None type in the given context. The type is owned by the
 /// context.
 MLIR_CAPI_EXPORTED MlirType mlirNoneTypeGet(MlirContext ctx);
 
 //===----------------------------------------------------------------------===//
 // Complex type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Complex type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirComplexTypeGetTypeID(void);
+
 /// Checks whether the given type is a Complex type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAComplex(MlirType type);
 
 /// Creates a complex type with the given element type in the same context as
 /// the element type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirComplexTypeGet(MlirType elementType);
 
 /// Returns the element type of the given complex type.
 MLIR_CAPI_EXPORTED MlirType mlirComplexTypeGetElementType(MlirType type);
 
 //===----------------------------------------------------------------------===//
 // Shaped type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Shaped type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirShapedTypeGetTypeID(void);
+
 /// Checks whether the given type is a Shaped type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAShaped(MlirType type);
 
 /// Returns the element type of the shaped type.
 MLIR_CAPI_EXPORTED MlirType mlirShapedTypeGetElementType(MlirType type);
 
 /// Checks whether the given shaped type is ranked.
 MLIR_CAPI_EXPORTED bool mlirShapedTypeHasRank(MlirType type);

 /// Prefer mlirShapedTypeGetDynamicStrideOrOffset to direct comparisons with
 /// this value.
 MLIR_CAPI_EXPORTED int64_t mlirShapedTypeGetDynamicStrideOrOffset(void);
 
 //===----------------------------------------------------------------------===//
 // Vector type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Vector type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirVectorTypeGetTypeID(void);
+
 /// Checks whether the given type is a Vector type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAVector(MlirType type);
 
 /// Creates a vector type of the shape identified by its rank and dimensions,
 /// with the given element type in the same context as the element type. The
 /// type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirVectorTypeGet(intptr_t rank,
                                               const int64_t *shape,
                                               MlirType elementType);
 
 /// Same as "mlirVectorTypeGet" but returns a nullptr wrapping MlirType on
 /// illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType mlirVectorTypeGetChecked(MlirLocation loc,
                                                      intptr_t rank,
                                                      const int64_t *shape,
                                                      MlirType elementType);
 
 //===----------------------------------------------------------------------===//
 // Ranked / Unranked Tensor type.
 //===----------------------------------------------------------------------===//
 
 /// Checks whether the given type is a Tensor type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsATensor(MlirType type);
 
+/// Returns the typeID of an RankedTensor type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirRankedTensorTypeGetTypeID(void);
+
 /// Checks whether the given type is a ranked tensor type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsARankedTensor(MlirType type);
 
+/// Returns the typeID of an UnrankedTensor type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirUnrankedTensorTypeGetTypeID(void);
+
 /// Checks whether the given type is an unranked tensor type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAUnrankedTensor(MlirType type);
 
 /// Creates a tensor type of a fixed rank with the given shape, element type,
 /// and optional encoding in the same context as the element type. The type is
 /// owned by the context. Tensor types without any specific encoding field
 /// should assign mlirAttributeGetNull() to this parameter.
 MLIR_CAPI_EXPORTED MlirType mlirRankedTensorTypeGet(intptr_t rank,

 /// on illegal arguments, emitting appropriate diagnostics.
 MLIR_CAPI_EXPORTED MlirType
 mlirUnrankedTensorTypeGetChecked(MlirLocation loc, MlirType elementType);
 
 //===----------------------------------------------------------------------===//
 // Ranked / Unranked MemRef type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an MemRef type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirMemRefTypeGetTypeID(void);
+
 /// Checks whether the given type is a MemRef type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAMemRef(MlirType type);
 
+/// Returns the typeID of an UnrankedMemRef type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirUnrankedMemRefTypeGetTypeID(void);
+
 /// Checks whether the given type is an UnrankedMemRef type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAUnrankedMemRef(MlirType type);
 
 /// Creates a MemRef type with the given rank and shape, a potentially empty
 /// list of affine layout maps, the given memory space and element type, in the
 /// same context as element type. The type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirMemRefTypeGet(MlirType elementType,
                                               intptr_t rank,

 /// Returns the memory spcae of the given Unranked MemRef type.
 MLIR_CAPI_EXPORTED MlirAttribute
 mlirUnrankedMemrefGetMemorySpace(MlirType type);
 
 //===----------------------------------------------------------------------===//
 // Tuple type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Tuple type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirTupleTypeGetTypeID(void);
+
 /// Checks whether the given type is a tuple type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsATuple(MlirType type);
 
 /// Creates a tuple type that consists of the given list of elemental types. The
 /// type is owned by the context.
 MLIR_CAPI_EXPORTED MlirType mlirTupleTypeGet(MlirContext ctx,
                                              intptr_t numElements,
                                              MlirType const *elements);
 
 /// Returns the number of types contained in a tuple.
 MLIR_CAPI_EXPORTED intptr_t mlirTupleTypeGetNumTypes(MlirType type);
 
 /// Returns the pos-th type in the tuple type.
 MLIR_CAPI_EXPORTED MlirType mlirTupleTypeGetType(MlirType type, intptr_t pos);
 
 //===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Function type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirFunctionTypeGetTypeID(void);
+
 /// Checks whether the given type is a function type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAFunction(MlirType type);
 
 /// Creates a function type, mapping a list of input types to result types.
 MLIR_CAPI_EXPORTED MlirType mlirFunctionTypeGet(MlirContext ctx,
                                                 intptr_t numInputs,
                                                 MlirType const *inputs,
                                                 intptr_t numResults,

 /// Returns the pos-th result type.
 MLIR_CAPI_EXPORTED MlirType mlirFunctionTypeGetResult(MlirType type,
                                                       intptr_t pos);
 
 //===----------------------------------------------------------------------===//
 // Opaque type.
 //===----------------------------------------------------------------------===//
 
+/// Returns the typeID of an Opaque type.
+MLIR_CAPI_EXPORTED MlirTypeID mlirOpaqueTypeGetTypeID(void);
+
 /// Checks whether the given type is an opaque type.
 MLIR_CAPI_EXPORTED bool mlirTypeIsAOpaque(MlirType type);
 
 /// Creates an opaque type in the given context associated with the dialect
 /// identified by its namespace. The type contains opaque byte data of the
 /// specified length (data need not be null-terminated).
 MLIR_CAPI_EXPORTED MlirType mlirOpaqueTypeGet(MlirContext ctx,
                                               MlirStringRef dialectNamespace,

mlir/include/mlir/Bindings/Python/PybindAdaptors.h

struct type_caster<MlirPassManager> {
   PYBIND11_TYPE_CASTER(MlirPassManager, _("MlirPassManager"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
     value = mlirPythonCapsuleToPassManager(capsule.ptr());
     return !mlirPassManagerIsNull(value);
   }
 };
 
+/// Casts object <-> MlirTypeID.
+template <>
+struct type_caster<MlirTypeID> {
+  PYBIND11_TYPE_CASTER(MlirTypeID, _("MlirTypeID"));
+  bool load(handle src, bool) {
+    py::object capsule = mlirApiObjectToCapsule(src);
+    value = mlirPythonCapsuleToTypeID(capsule.ptr());
+    return !mlirTypeIDIsNull(value);
+  }
+  static handle cast(MlirTypeID v, return_value_policy, handle) {
+    if (v.ptr == nullptr)
+      return py::none();
+    py::object capsule =
+        py::reinterpret_steal<py::object>(mlirPythonTypeIDToCapsule(v));
+    return py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+        .attr("TypeID")
+        .attr(MLIR_PYTHON_CAPI_FACTORY_ATTR)(capsule)
+        .release();
+  };
+};
+
 /// Casts object <-> MlirType.
 template <>
 struct type_caster<MlirType> {
   PYBIND11_TYPE_CASTER(MlirType, _("MlirType"));
   bool load(handle src, bool) {
     py::object capsule = mlirApiObjectToCapsule(src);
     value = mlirPythonCapsuleToType(capsule.ptr());
     return !mlirTypeIsNull(value);

mlir/lib/Bindings/Python/IRCore.cpp

 #include "PybindUtils.h"
 
 #include "mlir-c/Bindings/Python/Interop.h"
 #include "mlir-c/BuiltinAttributes.h"
 #include "mlir-c/Debug.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/Support.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/SmallVector.h"
 
 #include <optional>
 #include <utility>
 
 namespace py = pybind11;
 using namespace mlir;
PyType PyType::createFromCapsule(py::object capsule) {
   MlirType rawType = mlirPythonCapsuleToType(capsule.ptr());
   if (mlirTypeIsNull(rawType))
     throw py::error_already_set();
   return PyType(PyMlirContext::forContext(mlirTypeGetContext(rawType)),
                 rawType);
 }
 
 //------------------------------------------------------------------------------
+// PyTypeID.
+//------------------------------------------------------------------------------
+
+py::object PyTypeID::getCapsule() {
+  return py::reinterpret_steal<py::object>(mlirPythonTypeIDToCapsule(*this));
+}
+
+PyTypeID PyTypeID::createFromCapsule(py::object capsule) {
+  MlirTypeID mlirTypeID = mlirPythonCapsuleToTypeID(capsule.ptr());
+  if (mlirTypeIDIsNull(mlirTypeID))
+    throw py::error_already_set();
+  return PyTypeID(mlirTypeID);
+}
+bool PyTypeID::operator==(const PyTypeID &other) const {
+  return mlirTypeIDEqual(typeID, other.typeID);
+}
+
+//------------------------------------------------------------------------------
 // PyValue and subclases.
 //------------------------------------------------------------------------------
 
 pybind11::object PyValue::getCapsule() {
   return py::reinterpret_steal<py::object>(mlirPythonValueToCapsule(get()));
 }
 
 PyValue PyValue::createFromCapsule(pybind11::object capsule) {
py::class_<PyType>(m, "Type", py::module_local())
           "__str__",
           [](PyType &self) {
             PyPrintAccumulator printAccum;
             mlirTypePrint(self, printAccum.getCallback(),
                           printAccum.getUserData());
             return printAccum.join();
           },
           "Returns the assembly form of the type.")
-      .def("__repr__", [](PyType &self) {
+      .def("__repr__",
+           [](PyType &self) {
              // Generally, assembly formats are not printed for __repr__ because
              // this can cause exceptionally long debug output and exceptions.
              // However, types are an exception as they typically have compact
              // assembly forms and printing them is useful.
              PyPrintAccumulator printAccum;
              printAccum.parts.append("Type(");
-        mlirTypePrint(self, printAccum.getCallback(), printAccum.getUserData());
+             mlirTypePrint(self, printAccum.getCallback(),
+                           printAccum.getUserData());
              printAccum.parts.append(")");
              return printAccum.join();
+           })
+      .def_property_readonly("typeid", [](PyType &self) -> MlirTypeID {
+        MlirTypeID mlirTypeID = mlirTypeGetTypeID(self);
+        if (!mlirTypeIDIsNull(mlirTypeID))
+          return mlirTypeID;
+        auto origRepr =
+            pybind11::repr(pybind11::cast(self)).cast<std::string>();
+        throw py::value_error(
+            (origRepr + llvm::Twine(" has no typeid.")).str());
+      });
+
+  //----------------------------------------------------------------------------
+  // Mapping of PyTypeID.
+  //----------------------------------------------------------------------------
+  py::class_<PyTypeID>(m, "TypeID", py::module_local())
+      .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyTypeID::getCapsule)
+      .def(MLIR_PYTHON_CAPI_FACTORY_ATTR, &PyTypeID::createFromCapsule)
+      // Note, this tests whether the underlying TypeIDs are the same,
+      // not whether the wrapper MlirTypeIDs are the same, nor whether
+      // the Python objects are the same (i.e., PyTypeID is a value type).
+      .def("__eq__",
+           [](PyTypeID &self, PyTypeID &other) { return self == other; })
+      .def("__eq__",
+           [](PyTypeID &self, const py::object &other) { return false; })
+      // Note, this gives the hash value of the underlying TypeID, not the
+      // hash value of the Python object, nor the hash value of the
+      // MlirTypeID wrapper.
+      .def("__hash__", [](PyTypeID &self) {
+        return static_cast<size_t>(mlirTypeIDHashValue(self));
       });
 
   //----------------------------------------------------------------------------
   // Mapping of Value.
   //----------------------------------------------------------------------------
   py::class_<PyValue>(m, "Value", py::module_local())
       .def(py::init<PyValue &>(), py::keep_alive<0, 1>(), py::arg("value"))
       .def_property_readonly(MLIR_PYTHON_CAPI_PTR_ATTR, &PyValue::getCapsule)

mlir/lib/Bindings/Python/IRModule.h

 
 #include "PybindUtils.h"
 
 #include "mlir-c/AffineExpr.h"
 #include "mlir-c/AffineMap.h"
 #include "mlir-c/Diagnostics.h"
 #include "mlir-c/IR.h"
 #include "mlir-c/IntegerSet.h"
+#include "mlir/Bindings/Python/PybindAdaptors.h"
 #include "llvm/ADT/DenseMap.h"
 
 namespace mlir {
 namespace python {
 
 class PyBlock;
 class PyDiagnostic;
 class PyDiagnosticHandler;
public:
   /// may be a pre-existing object. Ownership of the underlying MlirType
   /// is taken by calling this function.
   static PyType createFromCapsule(pybind11::object capsule);
 
 private:
   MlirType type;
 };
 
+/// A TypeID provides an efficient and unique identifier for a specific C++
+/// type. This allows for a C++ type to be compared, hashed, and stored in an
+/// opaque context. This class wraps around the generic MlirTypeID.
+class PyTypeID {
+public:
+  PyTypeID(MlirTypeID typeID) : typeID(typeID) {}
+  // Note, this tests whether the underlying TypeIDs are the same,
+  // not whether the wrapper MlirTypeIDs are the same, nor whether
+  // the PyTypeID objects are the same (i.e., PyTypeID is a value type).
+  bool operator==(const PyTypeID &other) const;
+  operator MlirTypeID() const { return typeID; }
+  MlirTypeID get() { return typeID; }
+
+  /// Gets a capsule wrapping the void* within the MlirTypeID.
+  pybind11::object getCapsule();
+
+  /// Creates a PyTypeID from the MlirTypeID wrapped by a capsule.
+  static PyTypeID createFromCapsule(pybind11::object capsule);
+
+private:
+  MlirTypeID typeID;
+};
+
 /// CRTP base classes for Python types that subclass Type and should be
 /// castable from it (i.e. via something like IntegerType(t)).
 /// By default, type class hierarchies are one level deep (i.e. a
 /// concrete type class extends PyType); however, intermediate python-visible
 /// base classes can be modeled by specifying a BaseTy.
 template <typename DerivedTy, typename BaseTy = PyType>
 class PyConcreteType : public BaseTy {
 public:
   // Derived classes must define statics for:
   //   IsAFunctionTy isaFunction
   //   const char *pyClassName
   using ClassTy = pybind11::class_<DerivedTy, BaseTy>;
   using IsAFunctionTy = bool (*)(MlirType);
+  using GetTypeIDFunctionTy = MlirTypeID (*)();
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction = nullptr;
 
   PyConcreteType() = default;
   PyConcreteType(PyMlirContextRef contextRef, MlirType t)
-      : BaseTy(std::move(contextRef), t) {}
+      : BaseTy(std::move(contextRef), t) {
+    pybind11::implicitly_convertible<PyType, DerivedTy>();
+  }
   PyConcreteType(PyType &orig)
       : PyConcreteType(orig.getContext(), castFrom(orig)) {}
 
   static MlirType castFrom(PyType &orig) {
     if (!DerivedTy::isaFunction(orig)) {
       auto origRepr = pybind11::repr(pybind11::cast(orig)).cast<std::string>();
       throw SetPyError(PyExc_ValueError, llvm::Twine("Cannot cast type to ") +
                                              DerivedTy::pyClassName +
                                              " (from " + origRepr + ")");
     }
     return orig;
   }
 
   static void bind(pybind11::module &m) {
     auto cls = ClassTy(m, DerivedTy::pyClassName, pybind11::module_local());
     cls.def(pybind11::init<PyType &>(), pybind11::keep_alive<0, 1>(),
             pybind11::arg("cast_from_type"));
     cls.def_static(
         "isinstance",
         [](PyType &otherType) -> bool {
           return DerivedTy::isaFunction(otherType);
         },
         pybind11::arg("other"));
+    cls.def_property_readonly_static(
+        "static_typeid", [](py::object & /*class*/) -> MlirTypeID {
+          if (DerivedTy::getTypeIdFunction)
+            return DerivedTy::getTypeIdFunction();
+          throw SetPyError(PyExc_AttributeError,
+                           DerivedTy::pyClassName +
+                               llvm::Twine(" has no typeid."));
+        });
+    cls.def_property_readonly("typeid", [](PyType &self) {
+      return py::cast(self).attr("typeid").cast<MlirTypeID>();
+    });
+    cls.def("__repr__", [](DerivedTy &self) {
+      PyPrintAccumulator printAccum;
+      printAccum.parts.append(DerivedTy::pyClassName);
+      printAccum.parts.append("(");
+      mlirTypePrint(self, printAccum.getCallback(), printAccum.getUserData());
+      printAccum.parts.append(")");
+      return printAccum.join();
+    });
+
     DerivedTy::bindDerived(cls);
   }
 
   /// Implemented by derived classes to add methods to the Python subclass.
   static void bindDerived(ClassTy &m) {}
 };
 
 /// Wrapper around the generic MlirAttribute.

mlir/lib/Bindings/Python/IRTypes.cpp

 static int mlirTypeIsAIntegerOrFloat(MlirType type) {
   return mlirTypeIsAInteger(type) || mlirTypeIsABF16(type) ||
          mlirTypeIsAF16(type) || mlirTypeIsAF32(type) || mlirTypeIsAF64(type);
 }
 
 class PyIntegerType : public PyConcreteType<PyIntegerType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAInteger;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirIntegerTypeGetTypeID;
   static constexpr const char *pyClassName = "IntegerType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_signless",
         [](unsigned width, DefaultingPyMlirContext context) {
           MlirType t = mlirIntegerTypeGet(context->get(), width);
c.def_property_readonly(
         "Returns whether this is an unsigned integer");
   }
 };
 
 /// Index Type subclass - IndexType.
 class PyIndexType : public PyConcreteType<PyIndexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAIndex;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirIndexTypeGetTypeID;
   static constexpr const char *pyClassName = "IndexType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirIndexTypeGet(context->get());
           return PyIndexType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a index type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3FNType.
 class PyFloat8E4M3FNType : public PyConcreteType<PyFloat8E4M3FNType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FN;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E4M3FNTypeGetTypeID;
   static constexpr const char *pyClassName = "Float8E4M3FNType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3FNTypeGet(context->get());
           return PyFloat8E4M3FNType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3fn type.");
   }
 };
 
 /// Floating Point Type subclass - Float8M5E2Type.
 class PyFloat8E5M2Type : public PyConcreteType<PyFloat8E5M2Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E5M2TypeGetTypeID;
   static constexpr const char *pyClassName = "Float8E5M2Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E5M2TypeGet(context->get());
           return PyFloat8E5M2Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e5m2 type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3FNUZ.
 class PyFloat8E4M3FNUZType : public PyConcreteType<PyFloat8E4M3FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3FNUZ;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E4M3FNUZTypeGetTypeID;
   static constexpr const char *pyClassName = "Float8E4M3FNUZType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3FNUZTypeGet(context->get());
           return PyFloat8E4M3FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E4M3B11FNUZ.
 class PyFloat8E4M3B11FNUZType : public PyConcreteType<PyFloat8E4M3B11FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E4M3B11FNUZ;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E4M3B11FNUZTypeGetTypeID;
   static constexpr const char *pyClassName = "Float8E4M3B11FNUZType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E4M3B11FNUZTypeGet(context->get());
           return PyFloat8E4M3B11FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e4m3b11fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - Float8E5M2FNUZ.
 class PyFloat8E5M2FNUZType : public PyConcreteType<PyFloat8E5M2FNUZType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFloat8E5M2FNUZ;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat8E5M2FNUZTypeGetTypeID;
   static constexpr const char *pyClassName = "Float8E5M2FNUZType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirFloat8E5M2FNUZTypeGet(context->get());
           return PyFloat8E5M2FNUZType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a float8_e5m2fnuz type.");
   }
 };
 
 /// Floating Point Type subclass - BF16Type.
 class PyBF16Type : public PyConcreteType<PyBF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsABF16;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirBFloat16TypeGetTypeID;
   static constexpr const char *pyClassName = "BF16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirBF16TypeGet(context->get());
           return PyBF16Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a bf16 type.");
   }
 };
 
 /// Floating Point Type subclass - F16Type.
 class PyF16Type : public PyConcreteType<PyF16Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF16;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat16TypeGetTypeID;
   static constexpr const char *pyClassName = "F16Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF16TypeGet(context->get());
           return PyF16Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f16 type.");
   }
 };
 
 /// Floating Point Type subclass - F32Type.
 class PyF32Type : public PyConcreteType<PyF32Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF32;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat32TypeGetTypeID;
   static constexpr const char *pyClassName = "F32Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF32TypeGet(context->get());
           return PyF32Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f32 type.");
   }
 };
 
 /// Floating Point Type subclass - F64Type.
 class PyF64Type : public PyConcreteType<PyF64Type> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAF64;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFloat64TypeGetTypeID;
   static constexpr const char *pyClassName = "F64Type";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirF64TypeGet(context->get());
           return PyF64Type(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a f64 type.");
   }
 };
 
 /// None Type subclass - NoneType.
 class PyNoneType : public PyConcreteType<PyNoneType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsANone;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirNoneTypeGetTypeID;
   static constexpr const char *pyClassName = "NoneType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](DefaultingPyMlirContext context) {
           MlirType t = mlirNoneTypeGet(context->get());
           return PyNoneType(context->getRef(), t);
         },
         py::arg("context") = py::none(), "Create a none type.");
   }
 };
 
 /// Complex Type subclass - ComplexType.
 class PyComplexType : public PyConcreteType<PyComplexType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAComplex;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirComplexTypeGetTypeID;
   static constexpr const char *pyClassName = "ComplexType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType) {
           // The element must be a floating point or integer scalar type.
void requireHasRank() {
     }
   }
 };
 
 /// Vector Type subclass - VectorType.
 class PyVectorType : public PyConcreteType<PyVectorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAVector;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirVectorTypeGetTypeID;
   static constexpr const char *pyClassName = "VectorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            DefaultingPyLocation loc) {
public:
   }
 };
 
 /// Ranked Tensor Type subclass - RankedTensorType.
 class PyRankedTensorType
     : public PyConcreteType<PyRankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsARankedTensor;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirRankedTensorTypeGetTypeID;
   static constexpr const char *pyClassName = "RankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<int64_t> shape, PyType &elementType,
            std::optional<PyAttribute> &encodingAttr, DefaultingPyLocation loc) {
public:
   }
 };
 
 /// Unranked Tensor Type subclass - UnrankedTensorType.
 class PyUnrankedTensorType
     : public PyConcreteType<PyUnrankedTensorType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedTensor;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirUnrankedTensorTypeGetTypeID;
   static constexpr const char *pyClassName = "UnrankedTensorType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](PyType &elementType, DefaultingPyLocation loc) {
           PyMlirContext::ErrorCapture errors(loc->getContext());
           MlirType t = mlirUnrankedTensorTypeGetChecked(loc, elementType);
           if (mlirTypeIsNull(t))
             throw MLIRError("Invalid type", errors.take());
           return PyUnrankedTensorType(elementType.getContext(), t);
         },
         py::arg("element_type"), py::arg("loc") = py::none(),
         "Create a unranked tensor type");
   }
 };
 
 /// Ranked MemRef Type subclass - MemRefType.
 class PyMemRefType : public PyConcreteType<PyMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAMemRef;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirMemRefTypeGetTypeID;
   static constexpr const char *pyClassName = "MemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](std::vector<int64_t> shape, PyType &elementType,
             PyAttribute *layout, PyAttribute *memorySpace,
public:
   }
 };
 
 /// Unranked MemRef Type subclass - UnrankedMemRefType.
 class PyUnrankedMemRefType
     : public PyConcreteType<PyUnrankedMemRefType, PyShapedType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAUnrankedMemRef;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirUnrankedMemRefTypeGetTypeID;
   static constexpr const char *pyClassName = "UnrankedMemRefType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
          "get",
          [](PyType &elementType, PyAttribute *memorySpace,
             DefaultingPyLocation loc) {
c.def_static(
             "Returns the memory space of the given Unranked MemRef type.");
   }
 };
 
 /// Tuple Type subclass - TupleType.
 class PyTupleType : public PyConcreteType<PyTupleType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsATuple;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirTupleTypeGetTypeID;
   static constexpr const char *pyClassName = "TupleType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get_tuple",
         [](py::list elementList, DefaultingPyMlirContext context) {
           intptr_t num = py::len(elementList);
c.def_property_readonly(
         "Returns the number of types contained in a tuple.");
   }
 };
 
 /// Function type.
 class PyFunctionType : public PyConcreteType<PyFunctionType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAFunction;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirFunctionTypeGetTypeID;
   static constexpr const char *pyClassName = "FunctionType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::vector<PyType> inputs, std::vector<PyType> results,
            DefaultingPyMlirContext context) {

 static MlirStringRef toMlirStringRef(const std::string &s) {
   return mlirStringRefCreate(s.data(), s.size());
 }
 
 /// Opaque Type subclass - OpaqueType.
 class PyOpaqueType : public PyConcreteType<PyOpaqueType> {
 public:
   static constexpr IsAFunctionTy isaFunction = mlirTypeIsAOpaque;
+  static constexpr GetTypeIDFunctionTy getTypeIdFunction =
+      mlirOpaqueTypeGetTypeID;
   static constexpr const char *pyClassName = "OpaqueType";
   using PyConcreteType::PyConcreteType;
 
   static void bindDerived(ClassTy &c) {
     c.def_static(
         "get",
         [](std::string dialectNamespace, std::string typeData,
            DefaultingPyMlirContext context) {

mlir/lib/CAPI/IR/BuiltinTypes.cpp

 #include "mlir/IR/Types.h"
 
 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Integer types.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirIntegerTypeGetTypeID() { return wrap(IntegerType::getTypeID()); }
+
 bool mlirTypeIsAInteger(MlirType type) {
   return llvm::isa<IntegerType>(unwrap(type));
 }
 
 MlirType mlirIntegerTypeGet(MlirContext ctx, unsigned bitwidth) {
   return wrap(IntegerType::get(unwrap(ctx), bitwidth));
 }
 

 bool mlirIntegerTypeIsUnsigned(MlirType type) {
   return llvm::cast<IntegerType>(unwrap(type)).isUnsigned();
 }
 
 //===----------------------------------------------------------------------===//
 // Index type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirIndexTypeGetTypeID() { return wrap(IndexType::getTypeID()); }
+
 bool mlirTypeIsAIndex(MlirType type) {
   return llvm::isa<IndexType>(unwrap(type));
 }
 
 MlirType mlirIndexTypeGet(MlirContext ctx) {
   return wrap(IndexType::get(unwrap(ctx)));
 }
 
 //===----------------------------------------------------------------------===//
 // Floating-point types.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirFloat8E5M2TypeGetTypeID() {
+  return wrap(Float8E5M2Type::getTypeID());
+}
+
 bool mlirTypeIsAFloat8E5M2(MlirType type) {
   return unwrap(type).isFloat8E5M2();
 }
 
 MlirType mlirFloat8E5M2TypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E5M2(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat8E4M3FNTypeGetTypeID() {
+  return wrap(Float8E4M3FNType::getTypeID());
+}
+
 bool mlirTypeIsAFloat8E4M3FN(MlirType type) {
   return unwrap(type).isFloat8E4M3FN();
 }
 
 MlirType mlirFloat8E4M3FNTypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E4M3FN(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat8E5M2FNUZTypeGetTypeID() {
+  return wrap(Float8E5M2FNUZType::getTypeID());
+}
+
 bool mlirTypeIsAFloat8E5M2FNUZ(MlirType type) {
   return unwrap(type).isFloat8E5M2FNUZ();
 }
 
 MlirType mlirFloat8E5M2FNUZTypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E5M2FNUZ(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat8E4M3FNUZTypeGetTypeID() {
+  return wrap(Float8E4M3FNUZType::getTypeID());
+}
+
 bool mlirTypeIsAFloat8E4M3FNUZ(MlirType type) {
   return unwrap(type).isFloat8E4M3FNUZ();
 }
 
 MlirType mlirFloat8E4M3FNUZTypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E4M3FNUZ(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat8E4M3B11FNUZTypeGetTypeID() {
+  return wrap(Float8E4M3B11FNUZType::getTypeID());
+}
+
 bool mlirTypeIsAFloat8E4M3B11FNUZ(MlirType type) {
   return unwrap(type).isFloat8E4M3B11FNUZ();
 }
 
 MlirType mlirFloat8E4M3B11FNUZTypeGet(MlirContext ctx) {
   return wrap(FloatType::getFloat8E4M3B11FNUZ(unwrap(ctx)));
 }
 
+MlirTypeID mlirBFloat16TypeGetTypeID() {
+  return wrap(BFloat16Type::getTypeID());
+}
+
 bool mlirTypeIsABF16(MlirType type) { return unwrap(type).isBF16(); }
 
 MlirType mlirBF16TypeGet(MlirContext ctx) {
   return wrap(FloatType::getBF16(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat16TypeGetTypeID() { return wrap(Float16Type::getTypeID()); }
+
 bool mlirTypeIsAF16(MlirType type) { return unwrap(type).isF16(); }
 
 MlirType mlirF16TypeGet(MlirContext ctx) {
   return wrap(FloatType::getF16(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat32TypeGetTypeID() { return wrap(Float32Type::getTypeID()); }
+
 bool mlirTypeIsAF32(MlirType type) { return unwrap(type).isF32(); }
 
 MlirType mlirF32TypeGet(MlirContext ctx) {
   return wrap(FloatType::getF32(unwrap(ctx)));
 }
 
+MlirTypeID mlirFloat64TypeGetTypeID() { return wrap(Float64Type::getTypeID()); }
+
 bool mlirTypeIsAF64(MlirType type) { return unwrap(type).isF64(); }
 
 MlirType mlirF64TypeGet(MlirContext ctx) {
   return wrap(FloatType::getF64(unwrap(ctx)));
 }
 
 //===----------------------------------------------------------------------===//
 // None type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirNoneTypeGetTypeID() { return wrap(NoneType::getTypeID()); }
+
 bool mlirTypeIsANone(MlirType type) {
   return llvm::isa<NoneType>(unwrap(type));
 }
 
 MlirType mlirNoneTypeGet(MlirContext ctx) {
   return wrap(NoneType::get(unwrap(ctx)));
 }
 
 //===----------------------------------------------------------------------===//
 // Complex type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirComplexTypeGetTypeID() { return wrap(ComplexType::getTypeID()); }
+
 bool mlirTypeIsAComplex(MlirType type) {
   return llvm::isa<ComplexType>(unwrap(type));
 }
 
 MlirType mlirComplexTypeGet(MlirType elementType) {
   return wrap(ComplexType::get(unwrap(elementType)));
 }
 

 int64_t mlirShapedTypeGetDynamicStrideOrOffset() {
   return ShapedType::kDynamic;
 }
 
 //===----------------------------------------------------------------------===//
 // Vector type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirVectorTypeGetTypeID() { return wrap(VectorType::getTypeID()); }
+
 bool mlirTypeIsAVector(MlirType type) {
   return llvm::isa<VectorType>(unwrap(type));
 }
 
 MlirType mlirVectorTypeGet(intptr_t rank, const int64_t *shape,
                            MlirType elementType) {
   return wrap(VectorType::get(llvm::ArrayRef(shape, static_cast<size_t>(rank)),
                               unwrap(elementType)));

 //===----------------------------------------------------------------------===//
 // Ranked / Unranked tensor type.
 //===----------------------------------------------------------------------===//
 
 bool mlirTypeIsATensor(MlirType type) {
   return llvm::isa<TensorType>(unwrap(type));
 }
 
+MlirTypeID mlirRankedTensorTypeGetTypeID() {
+  return wrap(RankedTensorType::getTypeID());
+}
+
 bool mlirTypeIsARankedTensor(MlirType type) {
   return llvm::isa<RankedTensorType>(unwrap(type));
 }
 
+MlirTypeID mlirUnrankedTensorTypeGetTypeID() {
+  return wrap(UnrankedTensorType::getTypeID());
+}
+
 bool mlirTypeIsAUnrankedTensor(MlirType type) {
   return llvm::isa<UnrankedTensorType>(unwrap(type));
 }
 
 MlirType mlirRankedTensorTypeGet(intptr_t rank, const int64_t *shape,
                                  MlirType elementType, MlirAttribute encoding) {
   return wrap(
       RankedTensorType::get(llvm::ArrayRef(shape, static_cast<size_t>(rank)),
MlirType mlirUnrankedTensorTypeGetChecked(MlirLocation loc,
                                           MlirType elementType) {
   return wrap(UnrankedTensorType::getChecked(unwrap(loc), unwrap(elementType)));
 }
 
 //===----------------------------------------------------------------------===//
 // Ranked / Unranked MemRef type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirMemRefTypeGetTypeID() { return wrap(MemRefType::getTypeID()); }
+
 bool mlirTypeIsAMemRef(MlirType type) {
   return llvm::isa<MemRefType>(unwrap(type));
 }
 
 MlirType mlirMemRefTypeGet(MlirType elementType, intptr_t rank,
                            const int64_t *shape, MlirAttribute layout,
                            MlirAttribute memorySpace) {
   return wrap(MemRefType::get(

 MlirAffineMap mlirMemRefTypeGetAffineMap(MlirType type) {
   return wrap(llvm::cast<MemRefType>(unwrap(type)).getLayout().getAffineMap());
 }
 
 MlirAttribute mlirMemRefTypeGetMemorySpace(MlirType type) {
   return wrap(llvm::cast<MemRefType>(unwrap(type)).getMemorySpace());
 }
 
+MlirTypeID mlirUnrankedMemRefTypeGetTypeID() {
+  return wrap(UnrankedMemRefType::getTypeID());
+}
+
 bool mlirTypeIsAUnrankedMemRef(MlirType type) {
   return llvm::isa<UnrankedMemRefType>(unwrap(type));
 }
 
 MlirType mlirUnrankedMemRefTypeGet(MlirType elementType,
                                    MlirAttribute memorySpace) {
   return wrap(
       UnrankedMemRefType::get(unwrap(elementType), unwrap(memorySpace)));

 MlirAttribute mlirUnrankedMemrefGetMemorySpace(MlirType type) {
   return wrap(llvm::cast<UnrankedMemRefType>(unwrap(type)).getMemorySpace());
 }
 
 //===----------------------------------------------------------------------===//
 // Tuple type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirTupleTypeGetTypeID() { return wrap(TupleType::getTypeID()); }
+
 bool mlirTypeIsATuple(MlirType type) {
   return llvm::isa<TupleType>(unwrap(type));
 }
 
 MlirType mlirTupleTypeGet(MlirContext ctx, intptr_t numElements,
                           MlirType const *elements) {
   SmallVector<Type, 4> types;
   ArrayRef<Type> typeRef = unwrapList(numElements, elements, types);
   return wrap(TupleType::get(unwrap(ctx), typeRef));
 }
 
 intptr_t mlirTupleTypeGetNumTypes(MlirType type) {
   return llvm::cast<TupleType>(unwrap(type)).size();
 }
 
 MlirType mlirTupleTypeGetType(MlirType type, intptr_t pos) {
   return wrap(
       llvm::cast<TupleType>(unwrap(type)).getType(static_cast<size_t>(pos)));
 }
 
 //===----------------------------------------------------------------------===//
 // Function type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirFunctionTypeGetTypeID() {
+  return wrap(FunctionType::getTypeID());
+}
+
 bool mlirTypeIsAFunction(MlirType type) {
   return llvm::isa<FunctionType>(unwrap(type));
 }
 
 MlirType mlirFunctionTypeGet(MlirContext ctx, intptr_t numInputs,
                              MlirType const *inputs, intptr_t numResults,
                              MlirType const *results) {
   SmallVector<Type, 4> inputsList;
MlirType mlirFunctionTypeGetResult(MlirType type, intptr_t pos) {
   return wrap(llvm::cast<FunctionType>(unwrap(type))
                   .getResult(static_cast<unsigned>(pos)));
 }
 
 //===----------------------------------------------------------------------===//
 // Opaque type.
 //===----------------------------------------------------------------------===//
 
+MlirTypeID mlirOpaqueTypeGetTypeID() { return wrap(OpaqueType::getTypeID()); }
+
 bool mlirTypeIsAOpaque(MlirType type) {
   return llvm::isa<OpaqueType>(unwrap(type));
 }
 
 MlirType mlirOpaqueTypeGet(MlirContext ctx, MlirStringRef dialectNamespace,
                            MlirStringRef typeData) {
   return wrap(
       OpaqueType::get(StringAttr::get(unwrap(ctx), unwrap(dialectNamespace)),

mlir/python/mlir/dialects/python_test.py

 #  Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 #  See https://llvm.org/LICENSE.txt for license information.
 #  SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 from ._python_test_ops_gen import *
-from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue
+from .._mlir_libs._mlirPythonTest import TestAttr, TestType, TestTensorValue, TestTensorType
 
 def register_python_test_dialect(context, load=True):
   from .._mlir_libs import _mlirPythonTest
   _mlirPythonTest.register_python_test_dialect(context, load)

mlir/test/python/dialects/python_test.py

def testCustomType():
   with Context() as ctx:
     test.register_python_test_dialect(ctx)
     a = test.TestType.get()
     # CHECK: !python_test.test_type
     print(a)
 
     # The following cast must not assert.
     b = test.TestType(a)
+    # Instance custom types should have typeids
+    assert isinstance(b.typeid, TypeID)
+    # Subclasses of ir.Type should not have a static_typeid
+    # CHECK: 'TestType' object has no attribute 'static_typeid'
+    try:
+      b.static_typeid
+    except AttributeError as e:
+      print(e)
 
     i8 = IntegerType.get_signless(8)
     try:
       test.TestType(i8)
     except ValueError as e:
       assert "Cannot cast type to TestType" in str(e)
     else:
       raise
with InsertionPoint(module.body):
 
       tt = Tensor(t)
       # CHECK: Tensor(%{{.*}} = tensor.empty() : tensor<10x10xi8>)
       print(tt)
 
       # CHECK: False
       print(tt.is_null())
 
+      # Classes of custom types that inherit from concrete types should have
+      # static_typeid
+      assert isinstance(test.TestTensorType.static_typeid, TypeID)
+      # And it should be equal to the in-tree concrete type
+      assert test.TestTensorType.static_typeid == t.type.typeid
+
 
 # CHECK-LABEL: TEST: inferReturnTypeComponents
 @run
 def inferReturnTypeComponents():
     with Context() as ctx, Location.unknown(ctx):
         test.register_python_test_dialect(ctx)
         module = Module.create()
         i32 = IntegerType.get_signless(32)

mlir/test/python/ir/builtin_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
 
 
def testTypeHash():
 
   s = set()
   s.add(t1)
   s.add(t2)
   s.add(t3)
   # CHECK: len(s): 2
   print("len(s): ", len(s))
 
+
 # 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)
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))
 
+
 # CHECK-LABEL: TEST: testIndexType
 @run
 def testIndexType():
   with Context() as ctx:
     # CHECK: index type: index
     print("index type:", IndexType.get())
 
 
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)
+    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-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))
+    print("ranked tensor type:", RankedTensorType.get(shape, f32))
 
     none = NoneType.get()
     try:
       tensor_invalid = RankedTensorType.get(shape, none)
     except MLIRError as e:
       # CHECK: Invalid type:
       # CHECK: error: unknown: invalid tensor element type: 'none'
       print(e)
with Context() as ctx:
     # CHECK: pos-th type in the tuple type: f32
     print("pos-th type in the tuple type:", tuple_type.get_type(1))
 
 
 # CHECK-LABEL: TEST: testFunctionType
 @run
 def testFunctionType():
   with Context() as ctx:
-    input_types = [IntegerType.get_signless(32),
-                  IntegerType.get_signless(16)]
+    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)
 
 

 # CHECK-LABEL: TEST: testShapedTypeConstants
 # Tests that ShapedType exposes magic value constants.
 @run
 def testShapedTypeConstants():
   # CHECK: <class 'int'>
   print(type(ShapedType.get_dynamic_size()))
   # CHECK: <class 'int'>
   print(type(ShapedType.get_dynamic_stride_or_offset()))
+
+
+# CHECK-LABEL: TEST: testTypeIDs
+@run
+def testTypeIDs():
+  with Context(), Location.unknown():
+    f32 = F32Type.get()
+
+    types = [
+      (IntegerType, IntegerType.get_signless(16)),
+      (IndexType, IndexType.get()),
+      (Float8E4M3FNType, Float8E4M3FNType.get()),
+      (Float8E5M2Type, Float8E5M2Type.get()),
+      (Float8E4M3FNUZType, Float8E4M3FNUZType.get()),
+      (Float8E4M3B11FNUZType, Float8E4M3B11FNUZType.get()),
+      (Float8E5M2FNUZType, Float8E5M2FNUZType.get()),
+      (BF16Type, BF16Type.get()),
+      (F16Type, F16Type.get()),
+      (F32Type, F32Type.get()),
+      (F64Type, F64Type.get()),
+      (NoneType, NoneType.get()),
+      (ComplexType, ComplexType.get(f32)),
+      (VectorType, VectorType.get([2, 3], f32)),
+      (RankedTensorType, RankedTensorType.get([2, 3], f32)),
+      (UnrankedTensorType, UnrankedTensorType.get(f32)),
+      (MemRefType, MemRefType.get([2, 3], f32)),
+      (UnrankedMemRefType, UnrankedMemRefType.get(f32, Attribute.parse("2"))),
+      (TupleType, TupleType.get_tuple([f32])),
+      (FunctionType, FunctionType.get([], [])),
+      (OpaqueType, OpaqueType.get("tensor", "bob")),
+    ]
+
+    # CHECK: IntegerType(i16)
+    # CHECK: IndexType(index)
+    # CHECK: Float8E4M3FNType(f8E4M3FN)
+    # CHECK: Float8E5M2Type(f8E5M2)
+    # CHECK: Float8E4M3FNUZType(f8E4M3FNUZ)
+    # CHECK: Float8E4M3B11FNUZType(f8E4M3B11FNUZ)
+    # CHECK: Float8E5M2FNUZType(f8E5M2FNUZ)
+    # CHECK: BF16Type(bf16)
+    # CHECK: F16Type(f16)
+    # CHECK: F32Type(f32)
+    # CHECK: F64Type(f64)
+    # CHECK: NoneType(none)
+    # CHECK: ComplexType(complex<f32>)
+    # CHECK: VectorType(vector<2x3xf32>)
+    # CHECK: RankedTensorType(tensor<2x3xf32>)
+    # CHECK: UnrankedTensorType(tensor<*xf32>)
+    # CHECK: MemRefType(memref<2x3xf32>)
+    # CHECK: UnrankedMemRefType(memref<*xf32, 2>)
+    # CHECK: TupleType(tuple<f32>)
+    # CHECK: FunctionType(() -> ())
+    # CHECK: OpaqueType(!tensor.bob)
+    for _, t in types:
+      print(repr(t))
+
+    # Test getTypeIdFunction agrees with
+    # mlirTypeGetTypeID(self) for an instance.
+    # CHECK: all equal
+    for t1, t2 in types:
+      tid1, tid2 = t1.static_typeid, Type(t2).typeid
+      assert tid1 == tid2 and hash(tid1) == hash(
+          tid2), f"expected hash and value equality {t1} {t2}"
+    else:
+      print("all equal")
+
+    # Test that storing PyTypeID in python dicts
+    # works as expected.
+    typeid_dict = dict(types)
+    assert len(typeid_dict)
+
+    # CHECK: all equal
+    for t1, t2 in typeid_dict.items():
+      assert t1.static_typeid == t2.typeid and hash(
+          t1.static_typeid) == hash(
+              t2.typeid), f"expected hash and value equality {t1} {t2}"
+    else:
+      print("all equal")
+
+    # CHECK: ShapedType has no typeid.
+    try:
+      print(ShapedType.static_typeid)
+    except AttributeError as e:
+      print(e)
+
+    vector_type = Type.parse("vector<2x3xf32>")
+    # CHECK: True
+    print(ShapedType(vector_type).typeid == vector_type.typeid)

mlir/test/python/lib/PythonTestModule.cpp

 //===- PythonTestModule.cpp - Python extension for the PythonTest dialect -===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PythonTestCAPI.h"
+#include "mlir-c/BuiltinTypes.h"
 #include "mlir/Bindings/Python/PybindAdaptors.h"
 
 namespace py = pybind11;
 using namespace mlir::python::adaptors;
 
 PYBIND11_MODULE(_mlirPythonTest, m) {
   m.def(
       "register_python_test_dialect",
mlir_attribute_subclass(m, "TestAttr",
           py::arg("cls"), py::arg("context") = py::none());
   mlir_type_subclass(m, "TestType", mlirTypeIsAPythonTestTestType)
       .def_classmethod(
           "get",
           [](py::object cls, MlirContext ctx) {
             return cls(mlirPythonTestTestTypeGet(ctx));
           },
           py::arg("cls"), py::arg("context") = py::none());
+  mlir_type_subclass(m, "TestTensorType", mlirTypeIsARankedTensor,
+                     py::module::import(MAKE_MLIR_PYTHON_QUALNAME("ir"))
+                         .attr("RankedTensorType"));
   mlir_value_subclass(m, "TestTensorValue",
                       mlirTypeIsAPythonTestTestTensorValue)
       .def("is_null", [](MlirValue &self) { return mlirValueIsNull(self); });
 }