setup.py

from setuptools import setup, Command
from torch.utils import cpp_extension
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
import versioneer

import os

if hasattr(cpp_extension, "get_compiler_abi_compatibility_and_version"):
    # - New function since torch 1.12
    if cpp_extension.get_compiler_abi_compatibility_and_version("g++")[0]:
        os.environ["CC"] = "g++"
        os.environ["CXX"] = "g++"
    else:
        os.environ["CC"] = "c++"
        os.environ["CXX"] = "c++"
else:
    # - This works up to torch 1.11
    if cpp_extension.check_compiler_abi_compatibility("g++"):
        os.environ["CC"] = "g++"
        os.environ["CXX"] = "g++"
    else:
        os.environ["CC"] = "c++"
        os.environ["CXX"] = "c++"


# Class for clean command
class Cleaner(Command):
    """Clean command to tidy up after building"""

    user_options = []

    def initialize_options(self):
        pass

    def finalize_options(self):
        pass

    def run(self):
        os.system("rm -vrf ./build ./dist ./*pyc ./*egg-info")


# Update cmdclass
cmdclass = versioneer.get_cmdclass()
cmdclass.update({"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)})
cmdclass.update({"clean": Cleaner})

# Handle versions
version = versioneer.get_version()

# Install
setup(
    name="sinabs-exodus",
    version=version,
    packages=["sinabs.exodus", "sinabs.exodus.layers"],
    ext_modules=[
        CUDAExtension(
            name="exodus_cuda",
            sources=[
                "cuda/bindings.cu",
                # 'cuda/leaky_bindings.cu',
                # 'cuda/experimental_bindings.cu',
            ],
            depends=[
                "cuda/lif_kernels.h"
                "cuda/leaky_kernels.h"
                # 'cuda/experimental_kernels.h'
            ],
        )
    ],
    cmdclass=cmdclass,
    install_requires=["torch", f"sinabs >= 1.2.9"],
)