torch / torch
torch.triu_indices¶
-
torch.
triu_indices
(row, col, offset=0, *, dtype=torch.long, device='cpu', layout=torch.strided) → Tensor¶ Returns the indices of the upper triangular part of a
row
bycol
matrix in a 2-by-N Tensor, where the first row contains row coordinates of all indices and the second row contains column coordinates. Indices are ordered based on rows and then columns.The upper triangular part of the matrix is defined as the elements on and above the diagonal.
The argument
offset
controls which diagonal to consider. Ifoffset
= 0, all elements on and above the main diagonal are retained. A positive value excludes just as many diagonals above the main diagonal, and similarly a negative value includes just as many diagonals below the main diagonal. The main diagonal are the set of indices for where are the dimensions of the matrix.Note
When running on CUDA,
row * col
must be less than to prevent overflow during calculation.- Parameters
row (
int
) – number of rows in the 2-D matrix.col (
int
) – number of columns in the 2-D matrix.offset (
int
) – diagonal offset from the main diagonal. Default: if not provided, 0.
- Keyword Arguments
dtype (
torch.dtype
, optional) – the desired data type of returned tensor. Default: ifNone
,torch.long
.device (
torch.device
, optional) – the desired device of returned tensor. Default: ifNone
, uses the current device for the default tensor type (seetorch.set_default_tensor_type()
).device
will be the CPU for CPU tensor types and the current CUDA device for CUDA tensor types.layout (
torch.layout
, optional) – currently only supporttorch.strided
.
- Example::
>>> a = torch.triu_indices(3, 3) >>> a tensor([[0, 0, 0, 1, 1, 2], [0, 1, 2, 1, 2, 2]])
>>> a = torch.triu_indices(4, 3, -1) >>> a tensor([[0, 0, 0, 1, 1, 1, 2, 2, 3], [0, 1, 2, 0, 1, 2, 1, 2, 2]])
>>> a = torch.triu_indices(4, 3, 1) >>> a tensor([[0, 0, 1], [1, 2, 2]])
此页内容是否对您有帮助