Select

Select operation returns a tensor filled with the elements from the second or the third input, depending on the condition (the first input) value.

\[\dst[i] = cond[i] ? \src\_0[i] : \src\_1[i]\]

Broadcasting is supported.

If the auto_broadcast attribute is not none, the select operation takes a two-step broadcast before performing the selection:

• Step 1: Input tensors src_0 and src_1 are broadcasted to dst_shape according to the Numpy broadcast rules.

• Step 2: Then, the cond tensor will be one-way broadcasted to the dst_shape of broadcasted src_0 and src_1. To be more specific, we align the two shapes to the right and compare them from right to left. Each dimension should be either equal or the dimension of cond should be 1.

• example:

  • cond={4, 5}, dst_shape={2, 3, 4, 5} => dst = {2, 3, 4, 5}

  • cond={3, 1, 5}, dst_shape={2, 3, 4, 5} => dst = {2, 3, 4, 5}

  • cond={3,5}, dst_shape={2, 3, 4, 5} => dst = invalid_shape

Operation Attributes

Attribute Name	Description	Value Type	Supported Values	Required or Optional
auto_broadcast	Specifies rules used for auto-broadcasting of src tensors	string	none, numpy, (default)	Optional

Execution Arguments

The inputs and outputs must be provided according to the below index order when constructing an operation.

Inputs

Index	Argument Name	Required or Optional
0	cond	Required
1	src_0	Required
2	src_1	Required

@note All input shapes should match and no broadcasting is allowed if the auto_broadcast attribute is set to none, or can be broadcasted according to the broadcasting rules mentioned above if auto_broadcast attribute set to numpy.

Outputs

Index	Argument Name	Required or Optional
0	dst	Required

Supported Data Types

Select operation supports the following data type combinations.

Cond	Src_0	Src_1	Dst
boolean	f32	f32	f32
boolean	bf16	bf16	bf16
boolean	f16	f16	f16