ttnn.triu

ttnn.triu(input_tensor: ttnn.Tensor, *, diagonal: int = 0, memory_config: ttnn.MemoryConfig = None) → ttnn.Tensor

Returns the upper triangular part of input_tensor by zeroing out elements below the specified diagonal. Elements on and above the given diagonal are preserved, while elements below it are set to zero.

diagonal = 0 selects the main diagonal (keeps elements on and above the main diagonal)
diagonal > 0 selects a diagonal above the main diagonal (keeps fewer elements)
diagonal < 0 selects a diagonal below the main diagonal (keeps more elements)

\[\mathrm{{output\_tensor}}_i = \verb|triu|(\mathrm{{input\_tensor}}_i, \mathrm{diagonal})\]

Parameters:

input_tensor (ttnn.Tensor) – the input tensor.

Keyword Arguments:

diagonal (int, optional) – diagonal value. Defaults to 0.
memory_config (ttnn.MemoryConfig, optional) – memory configuration for the operation. Defaults to None.

Returns:

ttnn.Tensor – the output tensor.

Note

Supported dtypes and layouts:

Dtypes	Layouts
FLOAT32, BFLOAT16, BFLOAT8_B	TILE, ROW_MAJOR

Example

# Create a 4x4 tensor with values from 1 to 16
tensor = ttnn.from_torch(
    torch.arange(1, 17, dtype=torch.bfloat16).reshape(4, 4),
    dtype=ttnn.bfloat16,
    layout=ttnn.TILE_LAYOUT,
    device=device,
)

output_diag_0 = ttnn.triu(tensor, diagonal=0)
logger.info(f"Upper triangular (diagonal=0): {output_diag_0}")
# Result:
# [[ 1,  2,  3,  4],
#  [ 0,  6,  7,  8],
#  [ 0,  0, 11, 12],
#  [ 0,  0,  0, 16]]

output_diag_1 = ttnn.triu(tensor, diagonal=1)
logger.info(f"Upper triangular (diagonal=1): {output_diag_1}")
# Result:
# [[ 0,  2,  3,  4],
#  [ 0,  0,  7,  8],
#  [ 0,  0,  0, 12],
#  [ 0,  0,  0,  0]]

output_diag_neg1 = ttnn.triu(tensor, diagonal=-1)
logger.info(f"Upper triangular (diagonal=-1): {output_diag_neg1}")
# Result:
# [[ 1,  2,  3,  4],
#  [ 5,  6,  7,  8],
#  [ 0, 10, 11, 12],
#  [ 0,  0, 15, 16]]