TorchIntegral/setup.py at main · TheStageAI/TorchIntegral · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
import os
import platform
import numpy as np
from setuptools import setup
from setuptools.command.build_ext import build_ext
from distutils.extension import Extension
from Cython.Build import cythonize


# --------------------------------------------------------------------------------------------
# Handling compile args for Cython
# --------------------------------------------------------------------------------------------
if platform.system() == "Darwin":
    compile_opts = [
        "-std=c++11",
        "-Ofast",  # '-fopenmp',
        "-mmacosx-version-min={:}".format(platform.mac_ver()[0]),
    ]

elif platform.system() == "Linux":
    compile_opts = [
        "-std=c++11",
        "-Ofast",  # '-fopenmp'
    ]

else:
    raise EnvironmentError(
        "Not supported platform: {plat}".format(plat=platform.system())
    )


# --------------------------------------------------------------------------------------------
# C++/Cython extesnions and packages
# --------------------------------------------------------------------------------------------
tsp_ext = Extension(
    "torch_integral.tsp_solver.solver",
    sources=["torch_integral/tsp_solver/solver.pyx"],
    extra_compile_args=compile_opts,
    extra_link_args=compile_opts,
    language="c++",
    include_dirs=[np.get_include()],
)
ext_modules = [tsp_ext]
packages = [
    "torch_integral",
    "torch_integral.tsp_solver",
    "torch_integral.graph",
    "torch_integral.parametrizations",
]


# --------------------------------------------------------------------------------------------
# Package setup
# --------------------------------------------------------------------------------------------
setup(
    name="TorchIntegral",
    ext_modules=cythonize(ext_modules),
    version="0.0.0",
    author="Azim Kurbanov, Solodskikh Kirill",
    author_email="hello@thestage.ai",
    maintainer="TheStage.AI",
    maintainer_email="hello@thestage.ai",
    install_requires=["cython"],
    description="Official implementation of Integral Neural Networks in PyTorch.",
    url="https://inn.thestage.ai",
    zip_safe=False,
    packages=packages,
    license="Apache License 2.0",
    long_description="Bla Bla",
    classifiers=["Programming Language :: Python :: 3"],
)