setup.py 1.13 KB
Newer Older
rusty1s's avatar
rusty1s committed
1
import glob
rusty1s's avatar
rusty1s committed
2
3
from setuptools import setup, find_packages

rusty1s's avatar
rusty1s committed
4
5
6
7
8
9
10
11
12
13
14
15
16
import torch.cuda
from torch.utils.cpp_extension import CppExtension, CUDAExtension

ext_modules = [CppExtension('scatter_cpu', ['cpu/scatter.cpp'])]
cmdclass = {'build_ext': torch.utils.cpp_extension.BuildExtension}

if torch.cuda.is_available():
    ext_modules += [
        CUDAExtension('scatter_cuda',
                      ['cuda/scatter.cpp'] + glob.glob('cuda/*.cu'))
    ]

__version__ = '1.0.4'
rusty1s's avatar
rusty1s committed
17
url = 'https://github.com/rusty1s/pytorch_scatter'
rusty1s's avatar
rusty1s committed
18

rusty1s's avatar
rusty1s committed
19
install_requires = []
rusty1s's avatar
rusty1s committed
20
21
setup_requires = ['pytest-runner', 'cffi']
tests_require = ['pytest', 'pytest-cov']
rusty1s's avatar
rusty1s committed
22

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