where¶

nvtripy.where(condition: Tensor, input: Tensor | Number, other: Tensor | Number) → Tensor[source]¶

Returns a new tensor of elements selected from either input or other, depending on condition.

Parameters:

condition (Tensor) – [dtype=T2] The condition tensor. Where this is True, elements are selected from input. Otherwise, elements are selected from other.
input (Tensor | Number) – [dtype=T1] Tensor of values selected at indices where condition is True.
other (Tensor | Number) – [dtype=T1] Tensor values selected at indices where condition is False.

Returns:

[dtype=T1] A new tensor with the broadcasted shape.

Return type:

Tensor

Constraints:

All three parameters must be broadcast-compatible with each other.

DATA TYPE CONSTRAINTS:

T1: float16, bfloat16, float32, int8, int32, int64
T2: bool

Example

condition = tp.Tensor([[True, False], [True, True]])
input = tp.ones([2, 2], dtype=tp.float32)
other = tp.zeros([2, 2], dtype=tp.float32)
output = tp.where(condition, input, other)

Local Variables¶

>>> condition
tensor(
    [[True, False],
     [True, True]], 
    dtype=bool, loc=cpu:0, shape=(2, 2))

>>> input
tensor(
    [[1, 1],
     [1, 1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

>>> other
tensor(
    [[0, 0],
     [0, 0]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))

>>> output
tensor(
    [[1, 0],
     [1, 1]], 
    dtype=float32, loc=gpu:0, shape=(2, 2))