jactorch.functional.indexing#

Tensor indexing utils.

Functions

batched_index_select(tensor, batched_indices)

Select elements from tensor according to batched_indices.

index_nonzero(tensor, mask)

Iteratively generates the values of tensor where mask is nonzero.

index_one_hot(tensor, dim, index)

tensor[:, :, index, :]

index_one_hot_ellipsis(tensor, dim, index)

tensor[:, :, index, ...].

inverse_permutation(perm)

Inverse a permutation.

leftmost_nonzero(tensor, dim)

Return the smallest nonzero index along the dim axis.

one_hot(index, nr_classes)

Convert a list of class labels into one-hot representation.

one_hot_dim(index, nr_classes, dim)

Convert a tensor of class labels into one-hot representation by adding a new dimension indexed at dim.

one_hot_nd(index, nr_classes)

Convert a tensor of class labels into one-hot representation.

reversed(x[, dim])

Reverse a tensor along the given dimension.

rightmost_nonzero(tensor, dim)

Return the smallest nonzero index along the dim axis.

set_index_one_hot_(tensor, dim, index, value)

tensor[:, :, index, :, :] = value.

Functions

batched_index_select(tensor, batched_indices)[source]#

Select elements from tensor according to batched_indices. The first dimension is assumed to be the batch dimension. This operation is equivalent to numpy: tensor[np.arange(len(batched_indices)), batched_indices].

Parameters:
Returns:

the selected elements.

Return type:

torch.Tensor

index_nonzero(tensor, mask)[source]#

Iteratively generates the values of tensor where mask is nonzero. When mask is a 1D tensor, this function is equivalent to:

for i in range(mask.size(0)):
    if mask[i]:
        yield tensor[i]
Parameters:
Yields:

torch.Tensor – the values of tensor where mask is nonzero.

index_one_hot(tensor, dim, index)[source]#

tensor[:, :, index, :]

Parameters:

  • tensor (torch.Tensor) – input.

  • dim (int)

  • index (Tensor) – (torch.Tensor): the tensor containing the indices along the dim dimension.

Returns:

tensor[:, :, index, :, :].

Return type:

torch.Tensor

index_one_hot_ellipsis(tensor, dim, index)[source]#

tensor[:, :, index, …].

Parameters:

  • tensor (torch.Tensor) – input.

  • dim (int)

  • index (Tensor) – (torch.Tensor): the tensor containing the indices along the dim dimension.

Returns:

tensor[:, :, index, …].

Return type:

torch.Tensor

inverse_permutation(perm)[source]#

Inverse a permutation.

Warning

This function does not check the validness of the input. That is, if the input is not a permutation, this function may generate arbitrary output.

Parameters:

perm (torch.Tensor) – shape (N, ) representing a permutation of 0 ~ N - 1.

Returns:

the inverse permutation, which satisfies: inv[perm[x]] = x.

Return type:

torch.Tensor

leftmost_nonzero(tensor, dim)[source]#

Return the smallest nonzero index along the dim axis. The tensor should be binary.

Parameters:
Returns:

the smallest nonzero index along the dim axis.

Return type:

torch.Tensor

one_hot(index, nr_classes)[source]#

Convert a list of class labels into one-hot representation.

Note

This function support only one-dimensional input. For high dimensional inputs, use one_hot_nd.

Parameters:

  • index (torch.Tensor) – shape (N, ), input class labels.

  • nr_classes (int) – number of total classes.

Returns:

shape (N, nr_classes), one-hot representation of the class labels.

Return type:

torch.Tensor

one_hot_dim(index, nr_classes, dim)[source]#

Convert a tensor of class labels into one-hot representation by adding a new dimension indexed at dim.

Parameters:

  • index (torch.Tensor) – input class labels.

  • nr_classes (int) – number of total classes.

  • dim (int) – dimension of the class label.

Returns:

one-hot representation of the class labels.

Return type:

torch.Tensor

one_hot_nd(index, nr_classes)[source]#

Convert a tensor of class labels into one-hot representation.

Parameters:

  • index (torch.Tensor) – input class labels.

  • nr_classes (int) – number of total classes.

Returns:

one-hot representation of the class labels, the label dimension is assumed to be the last one.

Return type:

torch.Tensor

reversed(x, dim=-1)[source]#

Reverse a tensor along the given dimension. For example, if dim=0, it is equivalent to the python notation: x[::-1].

Parameters:
Returns:

of same shape as x, but with the dimension dim reversed.

Return type:

torch.Tensor

rightmost_nonzero(tensor, dim)[source]#

Return the smallest nonzero index along the dim axis. The tensor should be binary.

Parameters:
Returns:

the smallest nonzero index along the dim axis.

Return type:

torch.Tensor

set_index_one_hot_(tensor, dim, index, value)[source]#

tensor[:, :, index, :, :] = value.

Parameters:

  • tensor (torch.Tensor) – input.

  • dim (int)

  • index (Tensor) – (torch.Tensor): the tensor containing the indices along the dim dimension.

  • value (torch.Tensor) – the value to be set.

Return type:

None

On this page