tilus.Script.where¶
- Script.where(condition, x, y, *, out=None)[source]¶
Select elements from two tensors based on a condition.
This instruction selects elements from two tensors x and y based on a boolean condition tensor. For each element in the condition tensor, if the condition is True, the corresponding element from x is selected, otherwise the corresponding element from y is selected. The result is a new register tensor with the same dtype as x and y.
This instruction supports broadcasting, so if the shapes of x and y are not the same, they will be broadcasted to a common shape before performing the selection. The broadcasting rules are similar to those in NumPy and PyTorch, where dimensions of size 1 can be expanded to match the other tensor’s shape.
- Parameters:
condition (RegisterTensor) – The boolean register tensor that determines which elements to select from x and y. Its shape must match the shape of x and y after broadcasting.
x (RegisterTensor | Expr | int | float) – The register tensor or expression to select elements from when the condition is True. If x is not a register tensor, it will be converted to a zero-dimensional register tensor with the same dtype as y.
y (RegisterTensor | Expr | int | float) – The register tensor or expression to select elements from when the condition is False. If y is not a register tensor, it will be converted to a zero-dimensional register tensor with the same dtype as x.
out (RegisterTensor, optional) – The register tensor to store the result. If not provided, a new register tensor will be allocated.
- Returns:
ret – The register tensor containing the selected elements based on the condition. The shape of the output tensor will be determined by the broadcasting rules applied to x and y. The dtype of the output tensor will be the same as that of x and y.
- Return type: