"vscode:/vscode.git/clone" did not exist on "0799bc08431ebdf44884eccffe083896deb66828"
__init__.py 1.95 KB
Newer Older
rusty1s's avatar
update  
rusty1s committed
1
2
3
4
import importlib
import os.path as osp

import torch
rusty1s's avatar
rusty1s committed
5

Matthias Fey's avatar
Matthias Fey committed
6
__version__ = '1.6.1'
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
9
10
11
for library in [
        '_version', '_grid', '_graclus', '_fps', '_rw', '_sampler', '_nearest',
        '_knn', '_radius'
]:
rusty1s's avatar
rusty1s committed
12
13
14
15
16
17
18
19
20
21
    cuda_spec = importlib.machinery.PathFinder().find_spec(
        f'{library}_cuda', [osp.dirname(__file__)])
    cpu_spec = importlib.machinery.PathFinder().find_spec(
        f'{library}_cpu', [osp.dirname(__file__)])
    spec = cuda_spec or cpu_spec
    if spec is not None:
        torch.ops.load_library(spec.origin)
    else:  # pragma: no cover
        raise ImportError(f"Could not find module '{library}_cpu' in "
                          f"{osp.dirname(__file__)}")
rusty1s's avatar
update  
rusty1s committed
22

rusty1s's avatar
typo  
rusty1s committed
23
cuda_version = torch.ops.torch_cluster.cuda_version()
rusty1s's avatar
rusty1s committed
24
if torch.version.cuda is not None and cuda_version != -1:  # pragma: no cover
rusty1s's avatar
rusty1s committed
25
    if cuda_version < 10000:
rusty1s's avatar
update  
rusty1s committed
26
27
28
29
30
        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])
    t_major, t_minor = [int(x) for x in torch.version.cuda.split('.')]

rusty1s's avatar
rusty1s committed
31
    if t_major != major:
rusty1s's avatar
update  
rusty1s committed
32
33
34
35
36
37
38
        raise RuntimeError(
            f'Detected that PyTorch and torch_cluster were compiled with '
            f'different CUDA versions. PyTorch has CUDA version '
            f'{t_major}.{t_minor} and torch_cluster has CUDA version '
            f'{major}.{minor}. Please reinstall the torch_cluster that '
            f'matches your PyTorch install.')

rusty1s's avatar
rusty1s committed
39
from .fps import fps  # noqa
rusty1s's avatar
update  
rusty1s committed
40
41
from .graclus import graclus_cluster  # noqa
from .grid import grid_cluster  # noqa
rusty1s's avatar
rusty1s committed
42
from .knn import knn, knn_graph  # noqa
rusty1s's avatar
rusty1s committed
43
from .nearest import nearest  # noqa
rusty1s's avatar
rusty1s committed
44
from .radius import radius, radius_graph  # noqa
rusty1s's avatar
rusty1s committed
45
46
from .rw import random_walk  # noqa
from .sampler import neighbor_sampler  # noqa
rusty1s's avatar
rusty1s committed
47

rusty1s's avatar
rusty1s committed
48
49
50
51
__all__ = [
    'graclus_cluster',
    'grid_cluster',
    'fps',
rusty1s's avatar
rusty1s committed
52
53
54
55
56
    'nearest',
    'knn',
    'knn_graph',
    'radius',
    'radius_graph',
rusty1s's avatar
rusty1s committed
57
58
    'random_walk',
    'neighbor_sampler',
rusty1s's avatar
rusty1s committed
59
60
    '__version__',
]