where¶

dragon.where(
  inputs,
  **kwargs
)[source]¶

Select the elements from two branches under the condition.

\[\text{out}_{i} = \begin{cases} \text{input1}_{i}, & \text{ if } \text{condition}_{i} \\ \text{input2}_{i}, & \text{ otherwise } \end{cases} \]

Examples:

a = dragon.constant([1, 2, 3])
b = dragon.constant([3, 2, 1])
print(dragon.where([a > b, a, b]))  # [3, 2, 3]

If only the condition is given, return the coordinates of True elements:

x = dragon.constant([[True, False, True],
                     [False, True, True]])
print(dragon.where(x))  # [[0, 0], [0, 2], [1, 1], [1, 2]]

Parameters:

• inputs (Sequence[dragon.Tensor]) – The condition, input1 and input2 tensor.

Returns:

dragon.Tensor – The output tensor.

See also

dragon.nonzero(…)