Source code for torch_geometric.utils.sort

from typing import Optional, Tuple

import torch

import torch_geometric.typing
from torch_geometric.typing import pyg_lib

[docs]def index_sort( inputs: torch.Tensor, max_value: Optional[int] = None, ) -> Tuple[torch.Tensor, torch.Tensor]: r"""Sorts the elements of the :obj:`inputs` tensor in ascending order. It is expected that :obj:`inputs` is one-dimensional and that it only contains positive integer values. If :obj:`max_value` is given, it can be used by the underlying algorithm for better performance. Args: inputs (torch.Tensor): A vector with positive integer values. max_value (int, optional): The maximum value stored inside :obj:`inputs`. This value can be an estimation, but needs to be greater than or equal to the real maximum. (default: :obj:`None`) """ if not torch_geometric.typing.WITH_INDEX_SORT: # pragma: no cover return inputs.sort() return pyg_lib.ops.index_sort(inputs, max_value=max_value)