WhereOperation

Applicable Products

Hardware Model	Supported	Remarks
Atlas 350 accelerator cards	x	-
Atlas A3 inference products/Atlas A3 training products	√	-
Atlas A2 training products/Atlas A2 inference products	√	-
Atlas training products	√	-
Atlas inference products	√	The bf16 data type is not supported.
Atlas 200I/500 A2 inference products	√	-

Description

Ternary operation.

The input tensors are cond, x, and y, and the output tensor is z.

Input:
cond (condition tensor): defines the selection condition of each element.

x: first input tensor selected when the condition is 1.

y: second input tensor selected when the condition is 0.
Output: Each element of the output tensor is selected from x or y based on the condition of the corresponding position. A tensor with the same shape as x and y is returned.
z = (cond == 1 ? x : y)

Operator Context

Definition

struct WhereParam {
    uint8_t rsv[8] = {0};
};

Parameters

Member	Type	Default Value	Description
rsv[8]	uint8_t	{0}	Reserved

Input

Parameter	Dimension	Data Type	Format	Description
cond	[dim_0, dim_1, ..., dim_n]	int8	ND	Input tensor 1, which is a condition variable.
x	[x_dim_0, x_dim_1, ..., x_dim_n]	float16	ND	Input tensor 2
y	[y_dim_0, y_dim_1, ..., y_dim_n]	float16	ND	Input tensor 3

Output

Parameter	Dimension	Data Type	Format	Description
z	[z_dim_0, z_dim_1, ..., z_dim_n]	float16	ND	Result tensor

Constraints

The input cond element can only be 0 or 1. The dimension of output z is the broadcast result of inputs x and y. cond can be unidirectionally broadcast to both x and y.

API Calling Example

Input:

cond = [[1, 0], 
              [0, 1]]
x = [[1, 2],
       [3, 4]]
y = [[10, 20],
       [30, 40]]

Output:

z = [[ 1, 20],
       [30, 4]]

Parent topic: atb/infer_op_params.h