Construct a Call to allocate a storage with specific size, runtime_device_index, and dtype. Parameters ---------- shape : Expr The shape of the storage to be allocated. runtime_device_index : Union[int, Expr] The device index indicating on which device the tenso
(
shape: Expr,
runtime_device_index: int | Expr,
dtype: str | Expr,
storage_scope: str | StringImm = "global",
)
| 21 | |
| 22 | |
| 23 | def alloc_storage( |
| 24 | shape: Expr, |
| 25 | runtime_device_index: int | Expr, |
| 26 | dtype: str | Expr, |
| 27 | storage_scope: str | StringImm = "global", |
| 28 | ) -> Call: |
| 29 | """Construct a Call to allocate a storage with specific size, |
| 30 | runtime_device_index, and dtype. |
| 31 | |
| 32 | Parameters |
| 33 | ---------- |
| 34 | shape : Expr |
| 35 | The shape of the storage to be allocated. |
| 36 | |
| 37 | runtime_device_index : Union[int, Expr] |
| 38 | The device index indicating on which device the tensor is to |
| 39 | be allocated at runtime. Index -1 is reserved for the host device. |
| 40 | |
| 41 | dtype : Union[str, Expr] |
| 42 | The datatype of the storage to be allocated. |
| 43 | |
| 44 | storage_scope : Union[str, StringImm] |
| 45 | The storage scope of the storage to allocate. Default is global. |
| 46 | |
| 47 | Returns |
| 48 | ------- |
| 49 | result : Call |
| 50 | A relax Call, which gets the allocated storage. |
| 51 | """ |
| 52 | shape = convert_to_expr(shape) |
| 53 | if isinstance(dtype, str): |
| 54 | dtype = DataTypeImm(dtype) |
| 55 | if isinstance(storage_scope, str): |
| 56 | storage_scope = StringImm(storage_scope) |
| 57 | if isinstance(runtime_device_index, int): |
| 58 | runtime_device_index = PrimValue(runtime_device_index) |
| 59 | return _ffi_api.alloc_storage(shape, runtime_device_index, dtype, storage_scope) # type: ignore |
| 60 | |
| 61 | |
| 62 | def alloc_tensor( |
nothing calls this directly
no test coverage detected
searching dependent graphs…