Script.register_tensor¶
- Script.register_tensor(*, dtype, shape, init=None)¶
Create a register tensor.
This instruction allocates a register tensor with the specified data type, shape, (optional) layout, and (optional) initialization value.
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.
init (Callable[[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:
Notes
Thread group: Can be executed by any sized thread group.