Class DtypeConstructor

python/tvm/tirx/script/builder/ir.py:1528–1559

Callable + subscriptable dtype object. Replaces the plain functions previously returned by ``func_gen``. * ``T.float32()`` — same FFI call as before (returns ``Var``). * ``T.float32[N]`` — returns ``LocalVectorAnnotation("float32", (N,))``. * ``T.float32[M, N]`` — r

class DtypeConstructor:
    """Callable + subscriptable dtype object.

    Replaces the plain functions previously returned by ``func_gen``.

    * ``T.float32()``        — same FFI call as before (returns ``Var``).
    * ``T.float32[N]``       — returns ``LocalVectorAnnotation("float32", (N,))``.
    * ``T.float32[M, N]``    — returns ``LocalVectorAnnotation("float32", (M, N))``.
    * ``x: T.float32``       — parser calls this object, gets a ``Var``.
    """

    def __init__(self, ffi_name: str, dtype_str: str):
        self._ffi_name = ffi_name
        self._dtype_str = dtype_str

    def __call__(
        self,
        expr: "None | PrimExpr | Literal['inf', '-inf', 'nan'] | int | float" = None,
        *,
        is_size_var: bool = False,
    ) -> "PrimExpr":
        if isinstance(expr, str):
            expr = float(expr)
        return getattr(_ffi_api, self._ffi_name)(expr, is_size_var)

    def __getitem__(self, shape):
        if isinstance(shape, tuple):
            return LocalVectorAnnotation(self._dtype_str, shape)
        return LocalVectorAnnotation(self._dtype_str, (shape,))

    def __repr__(self):
        return f"DtypeConstructor({self._dtype_str!r})"


def allocate(