Construct a Call to allocate a tensor on a certain storage starting from the given offset. Parameters ---------- storage : Expr The storage to allocate the tensor to. offset : Union[int, Expr] The storage offset to allocate the tensor. shape : Expr The
(
storage: Expr,
offset: int | Expr,
shape: Expr,
dtype: str | Expr,
runtime_device_ind: int | Expr = PrimValue(0),
)
| 60 | |
| 61 | |
| 62 | def alloc_tensor( |
| 63 | storage: Expr, |
| 64 | offset: int | Expr, |
| 65 | shape: Expr, |
| 66 | dtype: str | Expr, |
| 67 | runtime_device_ind: int | Expr = PrimValue(0), |
| 68 | ) -> Call: |
| 69 | """Construct a Call to allocate a tensor on a certain storage starting from the given offset. |
| 70 | |
| 71 | Parameters |
| 72 | ---------- |
| 73 | storage : Expr |
| 74 | The storage to allocate the tensor to. |
| 75 | |
| 76 | offset : Union[int, Expr] |
| 77 | The storage offset to allocate the tensor. |
| 78 | |
| 79 | shape : Expr |
| 80 | The shape of the tensor to be allocated. |
| 81 | |
| 82 | dtype : Union[str, Expr] |
| 83 | The datatype of the tensor to be allocated. |
| 84 | |
| 85 | runtime_device_ind: Union[int, Expr] |
| 86 | The device index indicating on which device the tensor is to be |
| 87 | allocated at runtime. Index -1 is reserved for the host device. |
| 88 | |
| 89 | Returns |
| 90 | ------- |
| 91 | result : Call |
| 92 | A relax Call, which gets the allocated tensor. |
| 93 | """ |
| 94 | if isinstance(offset, int): |
| 95 | offset = PrimValue(offset) |
| 96 | shape = convert_to_expr(shape) |
| 97 | if isinstance(dtype, str): |
| 98 | dtype = DataTypeImm(dtype) |
| 99 | return _ffi_api.alloc_tensor(storage, offset, shape, dtype, runtime_device_ind) # type: ignore |
| 100 | |
| 101 | |
| 102 | def kill_storage(storage: Expr) -> Call: |
nothing calls this directly
no test coverage detected
searching dependent graphs…