where

→Tensor or tuple of Tensor

where(condition: Tensor, input: TensorOrScalar | None = ..., other: TensorOrScalar | None = ...)

source edit

Implementing kernel

C++where_opfree fn

Elementwise selection between input and other based on condition.

For each element position the result is input where condition is True and other otherwise. All three inputs broadcast against each other. where(cond) alone is nonzero(cond, as_tuple=True).

Parameters

conditionTensor

Boolean tensor.

inputTensor= ...

Values used where condition is True.

otherTensor= ...

Values used where condition is False.

Returns

Tensor or tuple of Tensor

Tensor of the broadcast shape, or with condition alone one int64 index tensor per dim.

Notes

Mathematically,

\text{out}_i \;=\; \begin{cases} x_i & \text{if } \text{cond}_i \\ y_i & \text{otherwise} . \end{cases}

Differentiable through both input and other — the gradient is gated by condition.

Examples

>>> import lucid
>>> a = lucid.tensor([1.0, 2.0, 3.0])
>>> b = lucid.tensor([-1.0, -2.0, -3.0])
>>> lucid.where(a > 1.5, a, b)
Tensor([-1.,  2.,  3.])