Function alloc_tensor

python/tvm/relax/op/builtin/builtin.py:23–65

Construct a Call to allocate a tensor with specific shape, dtype, runtime_device_index. Parameters ---------- shape : Expr The shape of the tensor to be allocated. dtype : Union[str, Expr] The datatype of the tensor to be allocated. runtime_device_index : Union

(
    shape: Expr,
    dtype: str | Expr,
    runtime_device_index: int | Expr,
    storage_scope: str | Expr = "global",
)

def alloc_tensor(
    shape: Expr,
    dtype: str | Expr,
    runtime_device_index: int | Expr,
    storage_scope: str | Expr = "global",
) -> Call:
    """Construct a Call to allocate a tensor with specific shape, dtype, runtime_device_index.

    Parameters
    ----------
    shape : Expr
        The shape of the tensor to be allocated.

    dtype : Union[str, Expr]
        The datatype of the tensor to be allocated.

    runtime_device_index : Union[int, Expr]
        The device index indicating on which device the tensor is to be allocated at runtime.
        Index -1 is reserved for the host device.

    storage_scope : Union[str, Expr]
        The storage scope to allocate the storage to.

    Returns
    -------
    result : Call
        A relax Call, which gets the allocated tensor.
    """
    if not isinstance(shape, Expr):
        shape = convert_to_expr(shape)
    if isinstance(dtype, str):
        dtype = DataTypeImm(dtype)
    if isinstance(runtime_device_index, int):
        runtime_device_index = PrimValue(runtime_device_index)
    if isinstance(storage_scope, str):
        storage_scope = StringImm(storage_scope)
    if not isinstance(storage_scope, StringImm):
        raise ValueError(
            "relax.builtin.alloc_tensor expects string as the storage scope, "
            f"but {storage_scope} is got."
        )

    return _ffi_api.alloc_tensor(shape, dtype, runtime_device_index, storage_scope)  # type: ignore


def stop_lift_params(x: Expr) -> Expr: