WhereOperation
Description
Ternary operation.
The input tensor includes cond, x, and y, and the output tensor is z (z = cond ? x : y).
Operator Context
Operator Function Implementation
- 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)
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 |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
int8 |
ND |
Output tensor 1, which is a condition variable. |
x |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16 |
ND |
Input tensor 2 |
y |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16 |
ND |
Input tensor 3 |
Output
Parameter |
Dimension |
Data Type |
Format |
Description |
|---|---|---|---|---|
z |
[-1,…,-1] The value -1 indicates that the size of the current dimension is not restricted. |
float16 |
ND |
Result tensor |
Restrictions
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, x, and y must be broadcastable.
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]]