tilus.Script.register_tensor¶
- Script.register_tensor(*, dtype, shape, layout=None, init=None)[source]¶
Create a register tensor.
This instruction allocates a register tensor with the specified data type, shape, (optional) layout, and (optional) initialization value.
When layout is not provided, the layout of the register tensor will be automatically inferred based on the operations performed on it.
If init is not provided, the register tensor will be uninitialized. If init is provided, it can be a scalar value (e.g., int, float, bool, or a scalar expression) that will be used to initialize all elements of the register tensor. It can also be a callable function that takes a sequence of index variables (i, j, …) and returns a scalar expression based on these indices. The element at index (i, j, …) will be initialized with the value returned by this function. When the data type of the value is not identical to the data type of the register tensor, the value will be cast to the data type of the register tensor.
- Parameters:
dtype (DataType) – The data type of the tensor elements.
shape (Sequence[int]) – The shape of the tensor.
layout (RegisterLayout, optional) – The layout of the tensor. If not provided, the layout will be inferred based on the operations performed on it.
init (Callable[[Sequence[Var]], Expr | int | float | bool] | Expr | int | float, optional) – The initialization value or function to initialize the tensor elements.
- Returns:
tensor – The allocated register tensor.
- Return type: