triton.where

triton.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) – IR builder to generate code into, optional from within @triton.jit functions