ttnn.where

ttnn.where(condition: 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 condition. For each element, if the corresponding entry in condition is 1, the output element is taken from true_value; otherwise, it is taken from false_value.

Parameters:

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

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

  • false_value (ttnn.Tensor or Number) – The value selected if the corresponding element in condition 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, layouts, and ranks:

Dtypes

Layouts

Ranks

BFLOAT16, BFLOAT8_B, FLOAT32, INT32

TILE

1, 2, 3, 4, 5

bfloat8_b/bfloat4_b supports only on TILE_LAYOUT

Example

# Create condition tensor and two value tensors
tensor1 = ttnn.from_torch(
    torch.tensor([[1, 2], [3, 4]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
tensor2 = ttnn.from_torch(
    torch.tensor([[5, 6], [7, 8]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)
tensor3 = ttnn.from_torch(
    torch.tensor([[9, 10], [11, 12]], dtype=torch.bfloat16), layout=ttnn.TILE_LAYOUT, device=device
)

# Perform the where operation
output = ttnn.where(tensor1, tensor2, tensor3)
logger.info(f"Where result: {output}")