axiom/pyproject.toml at main · VersesTech/axiom · 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
[build-system]
requires = ["setuptools>=61.0", "setuptools-scm>=8"]
build-backend = "setuptools.build_meta"

[project]
name = "axiom"
authors = [
    {name = "VERSES Research", email = "verses.research@verses.ai"},
]
description = "AXIOM: Learning to Play Games in Minutes with Expanding Object-Centric Models"
readme = "README.md"
requires-python = ">=3.10,<3.12"
dependencies = [
    "numpy==1.26.4",
    "jax==0.4.35",
    "opencv-python>=4.10.0",
    "wandb==0.18.3",
    "equinox==0.11.8",
    "optax==0.2.4",
    "mediapy==1.2.2",
    "moviepy==1.0.3",
    "rich==14.0.0",
    "multimethod>=1.12",
    "gameworld>=0.1.0",
]
version="0.1.0"

[tool.setuptools.packages.find]
include = ["axiom"]

[project.optional-dependencies]
gpu = [
    "jax[cuda12]==0.4.35",
    "nvidia-cublas-cu12==12.4.5.8",
    "nvidia-cuda-cupti-cu12==12.4.127",
    "nvidia-cuda-nvcc-cu12==12.8.93",
    "nvidia-cuda-nvrtc-cu12==12.4.127",
    "nvidia-cuda-runtime-cu12==12.4.127",
    "nvidia-cudnn-cu12==9.1.0.70",
    "nvidia-cufft-cu12==11.2.1.3",
    "nvidia-cufile-cu12==1.11.1.6",
    "nvidia-curand-cu12==10.3.5.147",
    "nvidia-cusolver-cu12==11.6.1.9",
    "nvidia-cusparse-cu12==12.3.1.170",
    "nvidia-cusparselt-cu12==0.6.2",
    "nvidia-nccl-cu12==2.21.5",
    "nvidia-nvjitlink-cu12==12.4.127",
    "nvidia-nvtx-cu12==12.4.127",
]