JaxKern/setup.py at main · JaxGaussianProcesses/JaxKern · 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 versioneer
from setuptools import find_packages, setup


readme = open("README.md").read()
NAME = "jaxkern"


# Handle builds of nightly release - adapted from BlackJax.
if "BUILD_JAXKERN_NIGHTLY" in os.environ:
    if os.environ["BUILD_JAXKERN_NIGHTLY"] == "nightly":
        NAME += "-nightly"
        from versioneer import get_versions as original_get_versions

        def get_versions():
            from datetime import datetime, timezone

            suffix = datetime.now(timezone.utc).strftime(r".dev%Y%m%d")
            versions = original_get_versions()
            versions["version"] = versions["version"].split("+")[0] + suffix
            return versions

        versioneer.get_versions = get_versions


REQUIRES = [
    "jax>=0.4.1",
    "jaxlib>=0.4.1",
    "jaxutils>=0.0.8",
    "jaxtyping>=0.0.2",
    "jaxlinop>=0.0.3",
    "deprecation",
    "distrax",
]

EXTRAS = {
    "dev": [
        "black",
        "isort",
        "pylint",
        "flake8",
        "pytest",
        "networkx",
        "pytest-cov",
        "pytest-xdist",
    ],
    "cuda": ["jax[cuda]"],
}


setup(
    name=NAME,
    version=versioneer.get_version(),
    cmdclass=versioneer.get_cmdclass(),
    author="Daniel Dodd and Thomas Pinder",
    author_email="tompinder@live.co.uk",
    packages=find_packages(".", exclude=["tests"]),
    license="LICENSE",
    description="Kernels in Jax.",
    long_description=readme,
    long_description_content_type="text/markdown",
    project_urls={
        "Documentation": "https://JaxKern.readthedocs.io/en/latest/",
        "Source": "https://github.com/JaxGaussianProcesses/JaxKern",
    },
    install_requires=REQUIRES,
    tests_require=EXTRAS["dev"],
    extras_require=EXTRAS,
    keywords=["gaussian-processes jax machine-learning bayesian"],
)