triu¶

nvtripy.triu(tensor: Tensor, diagonal: int = 0) → Tensor[source]¶

Returns the upper triangular part of each \([M, N]\) matrix in the tensor, with all other elements set to 0. If the tensor has more than two dimensions, it is treated as a batch of matrices.

Parameters:

tensor (Tensor) – The nvtripy tensor to operate on.
diagonal (int) –
The diagonal below which to zero elements. diagonal=0 indicates the main diagonal which is defined by the set of indices \({{(i, i)}}\) where \(i \in [0, min(M, N))\).

Positive values indicate the diagonal which is that many diagonals above the main one, while negative values indicate one which is below.

Returns:

A tensor of the same shape as this tensor.

Return type:

Tensor

Example: Main Diagonal

input = tp.iota((2, 1, 3, 3), dim=2) + 1.0
output = tp.triu(input)

Local Variables¶

>>> input
tensor(
    [[[[1.0000, 1.0000, 1.0000],
       [2.0000, 2.0000, 2.0000],
       [3.0000, 3.0000, 3.0000]]],


     [[[1.0000, 1.0000, 1.0000],
       [2.0000, 2.0000, 2.0000],
       [3.0000, 3.0000, 3.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(2, 1, 3, 3))

>>> output
tensor(
    [[[[1.0000, 1.0000, 1.0000],
       [0.0000, 2.0000, 2.0000],
       [0.0000, 0.0000, 3.0000]]],


     [[[1.0000, 1.0000, 1.0000],
       [0.0000, 2.0000, 2.0000],
       [0.0000, 0.0000, 3.0000]]]], 
    dtype=float32, loc=gpu:0, shape=(2, 1, 3, 3))

Example: Two Diagonals Above Main

output = tp.triu(input, diagonal=2)

Local Variables¶

>>> input
tensor(
    [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
     [2.0000, 2.0000, 2.0000, 2.0000, 2.0000],
     [3.0000, 3.0000, 3.0000, 3.0000, 3.0000],
     [4.0000, 4.0000, 4.0000, 4.0000, 4.0000],
     [5.0000, 5.0000, 5.0000, 5.0000, 5.0000]], 
    dtype=float32, loc=gpu:0, shape=(5, 5))

>>> output
tensor(
    [[0.0000, 0.0000, 1.0000, 1.0000, 1.0000],
     [0.0000, 0.0000, 0.0000, 2.0000, 2.0000],
     [0.0000, 0.0000, 0.0000, 0.0000, 3.0000],
     [0.0000, 0.0000, 0.0000, 0.0000, 0.0000],
     [0.0000, 0.0000, 0.0000, 0.0000, 0.0000]], 
    dtype=float32, loc=gpu:0, shape=(5, 5))

Example: One Diagonal Below Main

output = tp.triu(input, diagonal=-1)

Local Variables¶

>>> input
tensor(
    [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
     [2.0000, 2.0000, 2.0000, 2.0000, 2.0000],
     [3.0000, 3.0000, 3.0000, 3.0000, 3.0000],
     [4.0000, 4.0000, 4.0000, 4.0000, 4.0000],
     [5.0000, 5.0000, 5.0000, 5.0000, 5.0000]], 
    dtype=float32, loc=gpu:0, shape=(5, 5))

>>> output
tensor(
    [[1.0000, 1.0000, 1.0000, 1.0000, 1.0000],
     [2.0000, 2.0000, 2.0000, 2.0000, 2.0000],
     [0.0000, 3.0000, 3.0000, 3.0000, 3.0000],
     [0.0000, 0.0000, 4.0000, 4.0000, 4.0000],
     [0.0000, 0.0000, 0.0000, 5.0000, 5.0000]], 
    dtype=float32, loc=gpu:0, shape=(5, 5))