__init__.py 3.08 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.6'
rusty1s's avatar
rusty1s committed
7

rusty1s's avatar
rusty1s committed
8
9
for library in [
        '_version', '_convert', '_diag', '_spmm', '_spspmm', '_metis', '_rw',
rusty1s's avatar
rusty1s committed
10
        '_saint', '_sample', '_relabel'
rusty1s's avatar
rusty1s committed
11
12
13
]:
    torch.ops.load_library(importlib.machinery.PathFinder().find_spec(
        library, [osp.dirname(__file__)]).origin)
rusty1s's avatar
rusty1s committed
14

rusty1s's avatar
cleaner  
rusty1s committed
15
16
17
18
19
20
if torch.version.cuda is not None:  # pragma: no cover
    cuda_version = torch.ops.torch_sparse.cuda_version()

    if cuda_version == -1:
        major = minor = 0
    elif cuda_version < 10000:
rusty1s's avatar
rusty1s committed
21
22
23
24
25
26
27
28
29
30
31
32
33
        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('.')]

    if t_major != major or t_minor != minor:
        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
34
35
36
37
38
39
40
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
41
from .permute import permute  # noqa
rusty1s's avatar
typos  
rusty1s committed
42
43
44
45
46
from .diag import remove_diag, set_diag, fill_diag  # noqa
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
47
from .cat import cat  # noqa
rusty1s's avatar
rusty1s committed
48
from .rw import random_walk  # noqa
rusty1s's avatar
fix  
rusty1s committed
49
from .metis import partition  # noqa
rusty1s's avatar
rusty1s committed
50
from .bandwidth import reverse_cuthill_mckee  # noqa
rusty1s's avatar
rusty1s committed
51
from .saint import saint_subgraph  # noqa
rusty1s's avatar
rusty1s committed
52
from .padding import padded_index, padded_index_select  # noqa
rusty1s's avatar
rusty1s committed
53
from .sample import sample, sample_adj  # noqa
rusty1s's avatar
rusty1s committed
54

rusty1s's avatar
typos  
rusty1s committed
55
56
57
58
59
60
61
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
62
63

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