__init__.py

import importlib
import os.path as osp

import torch

__version__ = '1.2.1'

suffix = 'hip' if torch.cuda.is_available() else 'cpu'

for library in ['_version', '_basis', '_weighting']:
    torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
        f'{library}_{suffix}', [osp.dirname(__file__)]).origin)

if torch.cuda.is_available():  # pragma: no cover
    cuda_version = torch.ops.torch_spline_conv.cuda_version()

    if cuda_version == -1:
        major = minor = 0
    elif cuda_version < 10000:
        major, minor = int(str(cuda_version)[0]), int(str(cuda_version)[2])
    else:
        major, minor = int(str(cuda_version)[0:2]), int(str(cuda_version)[3])

from .basis import spline_basis  # noqa
from .weighting import spline_weighting  # noqa
from .conv import spline_conv  # noqa

__all__ = [
    'spline_basis',
    'spline_weighting',
    'spline_conv',
    '__version__',
]