__init__.py 2.23 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import importlib
rusty1s's avatar
typo  
rusty1s committed
2
3
4
5
import os.path as osp

import torch

rusty1s's avatar
rusty1s committed
6
7
from .scatter import (scatter_sum, scatter_add, scatter_mean, scatter_min,
                      scatter_max, scatter)
rusty1s's avatar
rusty1s committed
8
9
from .segment_csr import (segment_sum_csr, segment_add_csr, segment_mean_csr,
                          segment_min_csr, segment_max_csr, segment_csr,
rusty1s's avatar
rusty1s committed
10
                          gather_csr)
rusty1s's avatar
rusty1s committed
11
12
from .segment_coo import (segment_sum_coo, segment_add_coo, segment_mean_coo,
                          segment_min_coo, segment_max_coo, segment_coo,
rusty1s's avatar
rusty1s committed
13
                          gather_coo)
rusty1s's avatar
rusty1s committed
14
15
from .composite import (scatter_std, scatter_logsumexp, scatter_softmax,
                        scatter_log_softmax)
rusty1s's avatar
rusty1s committed
16

rusty1s's avatar
rusty1s committed
17
18
torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
    '_version', [osp.dirname(__file__)]).origin)
rusty1s's avatar
fix  
rusty1s committed
19
cuda_version = torch.ops.torch_scatter.cuda_version()
rusty1s's avatar
rusty1s committed
20

rusty1s's avatar
rusty1s committed
21
if cuda_version != -1 and torch.version.cuda is not None:  # pragma: no cover
rusty1s's avatar
rusty1s committed
22
23
24
25
26
    if 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])
    t_major, t_minor = [int(x) for x in torch.version.cuda.split('.')]
rusty1s's avatar
rusty1s committed
27
28
    cuda_version = str(major) + '.' + str(minor)

rusty1s's avatar
rusty1s committed
29
30
31
32
33
34
35
    if t_major != major or t_minor != minor:
        raise RuntimeError(
            'Detected that PyTorch and torch_scatter were compiled with '
            'different CUDA versions. PyTorch has CUDA version={}.{} and '
            'torch_scatter has CUDA version={}.{}. Please reinstall the '
            'torch_scatter that matches your PyTorch install.'.format(
                t_major, t_minor, major, minor))
rusty1s's avatar
rusty1s committed
36
37
else:
    cuda_version = None
rusty1s's avatar
typo  
rusty1s committed
38

rusty1s's avatar
update  
rusty1s committed
39
__version__ = '2.0.3'
rusty1s's avatar
cleaner  
rusty1s committed
40

rusty1s's avatar
rusty1s committed
41
__all__ = [
rusty1s's avatar
rusty1s committed
42
43
44
45
46
47
    'scatter_sum',
    'scatter_add',
    'scatter_mean',
    'scatter_min',
    'scatter_max',
    'scatter',
rusty1s's avatar
rusty1s committed
48
49
50
51
52
    'segment_sum_csr',
    'segment_add_csr',
    'segment_mean_csr',
    'segment_min_csr',
    'segment_max_csr',
rusty1s's avatar
rusty1s committed
53
    'segment_csr',
rusty1s's avatar
rusty1s committed
54
    'gather_csr',
rusty1s's avatar
rusty1s committed
55
56
57
58
59
60
61
    'segment_sum_coo',
    'segment_add_coo',
    'segment_mean_coo',
    'segment_min_coo',
    'segment_max_coo',
    'segment_coo',
    'gather_coo',
rusty1s's avatar
rusty1s committed
62
63
64
65
    'scatter_std',
    'scatter_logsumexp',
    'scatter_softmax',
    'scatter_log_softmax',
66
    'torch_scatter',
rusty1s's avatar
rusty1s committed
67
    'cuda_version',
rusty1s's avatar
rusty1s committed
68
    '__version__',
rusty1s's avatar
rusty1s committed
69
]