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

import torch

rusty1s's avatar
rusty1s committed
6
__version__ = '0.6.11'
rusty1s's avatar
rusty1s committed
7
8

suffix = 'cuda' if torch.cuda.is_available() else 'cpu'
rusty1s's avatar
rusty1s committed
9

rusty1s's avatar
rusty1s committed
10
for library in [
rusty1s's avatar
rusty1s committed
11
        '_version', '_convert', '_diag', '_spmm', '_spspmm', '_metis', '_rw',
rusty1s's avatar
rusty1s committed
12
        '_saint', '_sample', '_ego_sample', '_hgt_sample', '_relabel'
rusty1s's avatar
rusty1s committed
13
14
]:
    torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
rusty1s's avatar
rusty1s committed
15
        f'{library}_{suffix}', [osp.dirname(__file__)]).origin)
rusty1s's avatar
rusty1s committed
16

rusty1s's avatar
rusty1s committed
17
if torch.cuda.is_available():  # pragma: no cover
rusty1s's avatar
cleaner  
rusty1s committed
18
19
20
21
22
    cuda_version = torch.ops.torch_sparse.cuda_version()

    if cuda_version == -1:
        major = minor = 0
    elif cuda_version < 10000:
rusty1s's avatar
rusty1s committed
23
24
25
26
27
        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
28
    if t_major != major:
rusty1s's avatar
rusty1s committed
29
30
31
32
33
34
35
        raise RuntimeError(
            f'Detected that PyTorch and torch_sparse were compiled with '
            f'different CUDA versions. PyTorch has CUDA version '
            f'{t_major}.{t_minor} and torch_sparse has CUDA version '
            f'{major}.{minor}. Please reinstall the torch_sparse that '
            f'matches your PyTorch install.')

rusty1s's avatar
typos  
rusty1s committed
36
37
38
39
40
41
42
from .storage import SparseStorage  # noqa
from .tensor import SparseTensor  # noqa
from .transpose import t  # noqa
from .narrow import narrow, __narrow_diag__  # noqa
from .select import select  # noqa
from .index_select import index_select, index_select_nnz  # noqa
from .masked_select import masked_select, masked_select_nnz  # noqa
rusty1s's avatar
rusty1s committed
43
from .permute import permute  # noqa
rusty1s's avatar
rusty1s committed
44
from .diag import remove_diag, set_diag, fill_diag, get_diag  # noqa
rusty1s's avatar
typos  
rusty1s committed
45
46
47
48
from .add import add, add_, add_nnz, add_nnz_  # noqa
from .mul import mul, mul_, mul_nnz, mul_nnz_  # noqa
from .reduce import sum, mean, min, max  # noqa
from .matmul import matmul  # noqa
rusty1s's avatar
rusty1s committed
49
from .cat import cat  # noqa
rusty1s's avatar
rusty1s committed
50
from .rw import random_walk  # noqa
rusty1s's avatar
fix  
rusty1s committed
51
from .metis import partition  # noqa
rusty1s's avatar
rusty1s committed
52
from .bandwidth import reverse_cuthill_mckee  # noqa
rusty1s's avatar
rusty1s committed
53
from .saint import saint_subgraph  # noqa
rusty1s's avatar
rusty1s committed
54
from .padding import padded_index, padded_index_select  # noqa
rusty1s's avatar
rusty1s committed
55
from .sample import sample, sample_adj  # noqa
rusty1s's avatar
rusty1s committed
56

rusty1s's avatar
typos  
rusty1s committed
57
58
59
60
61
62
63
from .convert import to_torch_sparse, from_torch_sparse  # noqa
from .convert import to_scipy, from_scipy  # noqa
from .coalesce import coalesce  # noqa
from .transpose import transpose  # noqa
from .eye import eye  # noqa
from .spmm import spmm  # noqa
from .spspmm import spspmm  # noqa
rusty1s's avatar
rusty1s committed
64
65

__all__ = [
rusty1s's avatar
rusty1s committed
66
67
68
69
    'SparseStorage',
    'SparseTensor',
    't',
    'narrow',
rusty1s's avatar
rusty1s committed
70
    '__narrow_diag__',
rusty1s's avatar
rusty1s committed
71
72
73
74
75
    'select',
    'index_select',
    'index_select_nnz',
    'masked_select',
    'masked_select_nnz',
rusty1s's avatar
rusty1s committed
76
    'permute',
rusty1s's avatar
rusty1s committed
77
78
79
    'remove_diag',
    'set_diag',
    'fill_diag',
rusty1s's avatar
rusty1s committed
80
    'get_diag',
rusty1s's avatar
rusty1s committed
81
82
83
84
85
86
87
88
89
90
91
92
93
94
    'add',
    'add_',
    'add_nnz',
    'add_nnz_',
    'mul',
    'mul_',
    'mul_nnz',
    'mul_nnz_',
    'sum',
    'mean',
    'min',
    'max',
    'matmul',
    'cat',
rusty1s's avatar
rusty1s committed
95
    'random_walk',
rusty1s's avatar
fix  
rusty1s committed
96
    'partition',
rusty1s's avatar
rusty1s committed
97
    'reverse_cuthill_mckee',
rusty1s's avatar
rusty1s committed
98
    'saint_subgraph',
rusty1s's avatar
rusty1s committed
99
100
    'padded_index',
    'padded_index_select',
101
102
    'to_torch_sparse',
    'from_torch_sparse',
rusty1s's avatar
rusty1s committed
103
104
    'to_scipy',
    'from_scipy',
rusty1s's avatar
rusty1s committed
105
    'coalesce',
rusty1s's avatar
rusty1s committed
106
    'transpose',
rusty1s's avatar
rusty1s committed
107
    'eye',
rusty1s's avatar
rusty1s committed
108
    'spmm',
rusty1s's avatar
rusty1s committed
109
    'spspmm',
rusty1s's avatar
rusty1s committed
110
    '__version__',
rusty1s's avatar
rusty1s committed
111
]