NeuralTangentEnsemble/pyproject.toml at main · ZadorLaboratory/NeuralTangentEnsemble · 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
[tool.poetry]
name = "neural-tangent-ensemble"
version = "0.1.0"
description = ""
authors = ["Ari Benjamin <benjami@cshl.edu>"]
license = "MIT"
readme = "README.md"

[tool.poetry.dependencies]
python = "~3.11"
ipython = "*"
tensorflow = "~2.15"
tensorflow_datasets = "*"
nvidia-cudnn-cu12 = "8.9.4.*"
jax = {version="~0.4", extras=["cuda12_pip"]}
jaxlib = {version="~0.4", source="jaxcuda"}
optax = "*"
orbax-checkpoint = "*"
flax = "*"
clu = "*"
hydra-core = "~1.3"
wandb = "*"
seaborn = "*"
nvidia-cublas-cu12 = "^12.5.3.2"

[[tool.poetry.source]]
name = "PyPI"
priority = "primary"

[[tool.poetry.source]]
name = "jaxcuda"
url = "https://storage.googleapis.com/jax-releases/jax_cuda_releases.html"
priority = "explicit"

[build-system]
requires = ["poetry-core"]
build-backend = "poetry.core.masonry.api"