Function alloc_storage

python/tvm/relax/op/memory/memory.py:23–59

Construct a Call to allocate a storage with specific size, virtual_device_index, storage_scope and dtype. Parameters ---------- size : Expr The size of the storage to be allocated. virtual_device_index : Union[int, Expr] The virtual device index indicating on wh

(
    size: Expr,
    virtual_device_index: int | Expr,
    storage_scope: str | Expr,
    dtype: str | Expr,
)

def alloc_storage(
    size: Expr,
    virtual_device_index: int | Expr,
    storage_scope: str | Expr,
    dtype: str | Expr,
) -> Call:
    """Construct a Call to allocate a storage with specific size, virtual_device_index,
    storage_scope and dtype.

    Parameters
    ----------
    size : Expr
        The size of the storage to be allocated.

    virtual_device_index : Union[int, Expr]
        The virtual device index indicating on which device the storage is to be allocated.
        Index -1 is reserved for the host device.

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

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

    Returns
    -------
    result : Call
        A relax Call, which gets the allocated storage.
    """
    size = convert_to_expr(size)
    if isinstance(dtype, str):
        dtype = DataTypeImm(dtype)
    if isinstance(storage_scope, str):
        storage_scope = StringImm(storage_scope)
    if isinstance(virtual_device_index, int):
        virtual_device_index = PrimValue(virtual_device_index)
    return _ffi_api.alloc_storage(size, virtual_device_index, storage_scope, dtype)  # type: ignore


def alloc_tensor(