utils.py

from typing import Any, Optional, Tuple

import torch

import torch_sparse.typing
from torch_sparse.typing import pyg_lib

try:
    from typing_extensions import Final  # noqa
except ImportError:
    from torch.jit import Final  # noqa


def index_sort(
        inputs: torch.Tensor,
        max_value: Optional[int] = None) -> Tuple[torch.Tensor, torch.Tensor]:
    r"""See pyg-lib documentation for more details:
    https://pyg-lib.readthedocs.io/en/latest/modules/ops.html"""
    if not torch_sparse.typing.WITH_INDEX_SORT:  # pragma: no cover
        return inputs.sort()
    return pyg_lib.ops.index_sort(inputs, max_value)


def is_scalar(other: Any) -> bool:
    return isinstance(other, int) or isinstance(other, float)