ttnn.where

ttnn.where(predicate: ttnn.Tensor, true_value: ttnn.Tensor or Number, false_value: ttnn.Tensor or Number, *, memory_config: ttnn.MemoryConfig = None, output_tensor: ttnn.Tensor = None, sub_core_grids: ttnn.CoreRangeSet = None) → None

Selects elements from true_value or false_value depending on the corresponding value in predicate. For each element, if the corresponding entry in predicate is 1, the output element is taken from true_value; otherwise, it is taken from false_value.

Parameters:

  • predicate (ttnn.Tensor)the predicate tensor must contain only 0’s or 1’s.

  • true_value (ttnn.Tensor or Number)The value selected if the corresponding element in predicate is 1.

  • false_value (ttnn.Tensor or Number)The value selected if the corresponding element in predicate is 0.

Keyword Arguments:

  • memory_config (ttnn.MemoryConfig, optional)memory configuration for the operation. Defaults to None.

  • output_tensor (ttnn.Tensor, optional)preallocated output tensor. Defaults to None.

  • sub_core_grids (ttnn.CoreRangeSet, optional)sub core grids for the operation. Defaults to None.

Note

Supported dtypes and layouts:

Dtypes

Layouts

BFLOAT16, BFLOAT8_B, BFLOAT4_B, FLOAT32, INT32, UINT32 (range: [0, 4294967295])

TILE

bfloat8_b/bfloat4_b supports only on TILE_LAYOUT

Example

# Create predicate tensor of 0's and 1's and two value tensors
predicate = ttnn.from_torch(
    torch.tensor([[1, 0], [0, 1]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
true_value = ttnn.from_torch(
    torch.tensor([[5, 6], [7, 8]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
false_value = ttnn.from_torch(
    torch.tensor([[9, 10], [11, 12]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)

# Perform the where operation, giving [[5, 10], [11, 8]]
output = ttnn.where(predicate, true_value, false_value)
logger.info(f"Where result: {output}")