asc2.where

asc2.where(mask: LocalTensor, src0: LocalTensor | PlainValue | int | float, src1: LocalTensor | PlainValue | int | float) → LocalTensor

Select elements from two sources based on a mask.

For each element, returns the corresponding element from src0 if the mask element is true (non-zero), otherwise returns the element from src1.

The supported data types for src0 and src1: int16, int32, float16, bfloat16, float32.

Parameters:

• mask – A boolean tensor specifying which elements to select

• src0 – The source for elements where mask is true (tensor or scalar)

• src1 – The source for elements where mask is false (tensor or scalar)

Returns:

A tensor with elements selected from src0 or src1 based on the mask

Return type:

LocalTensor

Raises:

• TypeError – If mask is not a LocalTensor, or if src0 or src1 is not a LocalTensor or scalar

• RuntimeError – If mask dtype is not int1, or if src0 or src1 dtype is not supported

Note

At least one of src0 or src1 must be a tensor with the same shape as the mask. Scalars are broadcast to the mask shape.

Examples

Select elements from two tensors based on a mask:

mask = tensor_a > tensor_b
result = asc2.where(mask, tensor_a, tensor_b)

Select elements from a tensor or a scalar based on a mask:

mask = tensor > 0
result = asc2.where(mask, tensor, 0)