setup.py 1.07 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
2
import torch
from setuptools import setup, find_packages
rusty1s's avatar
rusty1s committed
3
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
rusty1s's avatar
rusty1s committed
4

rusty1s's avatar
rusty1s committed
5
__version__ = '0.2.0'
rusty1s's avatar
rusty1s committed
6
7
8
9
10
11
12
13
14
url = 'https://github.com/rusty1s/pytorch_sparse'

install_requires = ['numpy', 'scipy']
setup_requires = ['pytest-runner']
tests_require = ['pytest', 'pytest-cov']
ext_modules = []
cmdclass = {}

if torch.cuda.is_available():
rusty1s's avatar
to csr  
rusty1s committed
15
16
    ext_modules += [
        CUDAExtension('matmul_cuda',
rusty1s's avatar
rename  
rusty1s committed
17
                      ['cuda/matmul.cpp', 'cuda/matmul_kernel.cu'])
rusty1s's avatar
to csr  
rusty1s committed
18
    ]
rusty1s's avatar
rusty1s committed
19
    cmdclass['build_ext'] = BuildExtension
rusty1s's avatar
rusty1s committed
20
21
22
23

setup(
    name='torch_sparse',
    version=__version__,
rusty1s's avatar
rusty1s committed
24
25
    description='PyTorch Extension Library of Optimized Sparse Matrix '
    'Operations',
rusty1s's avatar
rusty1s committed
26
27
28
29
30
31
32
33
34
35
36
37
    author='Matthias Fey',
    author_email='matthias.fey@tu-dortmund.de',
    url=url,
    download_url='{}/archive/{}.tar.gz'.format(url, __version__),
    keywords=['pytorch', 'sparse', 'deep-learning'],
    install_requires=install_requires,
    setup_requires=setup_requires,
    tests_require=tests_require,
    ext_modules=ext_modules,
    cmdclass=cmdclass,
    packages=find_packages(),
)