Register object type. Parameters ---------- type_key The type key of the node. It requires ``type_key`` to be registered already on the C++ side. If not specified, the class name will be used. init If True (default), install ``__init__`` from the C++ ``__ffi_
(
type_key: str | None = None,
*,
init: bool = True,
)
| 34 | |
| 35 | |
| 36 | def register_object( |
| 37 | type_key: str | None = None, |
| 38 | *, |
| 39 | init: bool = True, |
| 40 | ) -> Callable[[_T], _T]: |
| 41 | """Register object type. |
| 42 | |
| 43 | Parameters |
| 44 | ---------- |
| 45 | type_key |
| 46 | The type key of the node. It requires ``type_key`` to be registered already |
| 47 | on the C++ side. If not specified, the class name will be used. |
| 48 | init |
| 49 | If True (default), install ``__init__`` from the C++ ``__ffi_init__`` |
| 50 | TypeAttrColumn when available, or a TypeError guard for ``Object`` |
| 51 | subclasses that lack one. Set to False when a subsequent decorator |
| 52 | (e.g. ``@c_class``) will handle ``__init__`` installation. |
| 53 | |
| 54 | Notes |
| 55 | ----- |
| 56 | All :class:`Object` subclasses get ``__slots__ = ()`` by default via the |
| 57 | metaclass, preventing per-instance ``__dict__``. To opt out and allow |
| 58 | arbitrary instance attributes, declare ``__slots__ = ("__dict__",)`` |
| 59 | explicitly in the class body:: |
| 60 | |
| 61 | @tvm_ffi.register_object("test.MyObject") |
| 62 | class MyObject(Object): |
| 63 | __slots__ = ("__dict__",) |
| 64 | |
| 65 | Examples |
| 66 | -------- |
| 67 | The following code registers MyObject using type key "test.MyObject", if the |
| 68 | type key is already registered on the C++ side. |
| 69 | |
| 70 | .. code-block:: python |
| 71 | |
| 72 | @tvm_ffi.register_object("test.MyObject") |
| 73 | class MyObject(Object): |
| 74 | pass |
| 75 | |
| 76 | """ |
| 77 | |
| 78 | def _register(cls: _T, object_name: str) -> _T: |
| 79 | """Register the object type with the FFI core.""" |
| 80 | type_index = core._object_type_key_to_index(object_name) |
| 81 | if type_index is None: |
| 82 | if _SKIP_UNKNOWN_OBJECTS: |
| 83 | return cls |
| 84 | raise ValueError(f"Cannot find object type index for {object_name}") |
| 85 | info = core._register_object_by_index(type_index, cls) |
| 86 | _add_class_attrs(type_cls=cls, type_info=info) |
| 87 | setattr(cls, "__tvm_ffi_type_info__", info) |
| 88 | if init: |
| 89 | _install_init(cls, info) |
| 90 | return cls |
| 91 | |
| 92 | if isinstance(type_key, str): |
| 93 |
no test coverage detected