triton.language.where

triton.language.where(condition, x, y, builder=None)

Returns a block of elements from either x or y, depending on condition.

Note that x and y are always evaluated regardless of the value of condition.

If you want to avoid unintented memory operations, use the mask arguments in triton.load and triton.store instead.

The shape of x and y are both broadcast to the shape of condition. x and y must have the data type.

Parameters

• condition (Block of triton.bool) – When True (nonzero), yield x, otherwise yield y.

• x – values selected at indices where condition is True.

• y – values selected at indices where condition is False.

• builder (triton.ir.builder, optional from within JIT'ed functions) – IR builder to generate code into