pyproject.toml

[tool.poetry]
name = "MILE"
version = "0.1.0"
description = "Implementation of Microcanonical Langevin Ensembles (MILE)."
authors = []
readme = "README.md"
package-mode = false

[tool.poetry.dependencies]
python = ">=3.10"
jax = "0.4.28"
blackjax = "1.2.2"
optax = "^0.2.1"
orbax = "^0.1.9"
numpyro = "^0.14.0"
flax = "^0.8.2"
matplotlib = "^3.8.0"
jupyter = "1.0.0"
pandas = "2.1.1"
typer = "0.9.0"
seaborn = "^0.13.0"
clu = "0.0.9"
nbconvert = "^7.16.4"
numpy = "1.26.4"
scikit-learn = "^1.5.2"

[tool.poetry.group.cuda]
optional = true

[tool.poetry.group.cuda.dependencies]
jax = {version="0.4.28", extras=["cuda12"]}

[tool.poetry.group.vision]
optional = true
[tool.poetry.group.vision.dependencies]
torch = [
    {version = "^2.2.1", platform = "linux", source = "pytorch-cpu"},
    {version = "^2.2.1", platform = "!=linux", source = "PyPI"}
]
torchvision = [
    {version = "^0.17.1", platform = "linux", source = "pytorch-cpu"},
    {version = "^0.17.1", platform = "!=linux", source = "PyPI"}
]

[tool.poetry.group.nlp]
optional = true
[tool.poetry.group.nlp.dependencies]
tokenizers = "^0.19.1"
datasets = "^2.19.0"
tiktoken = "^0.6.0"

[[tool.poetry.source]]
name = "pytorch-cpu"
url = "https://download.pytorch.org/whl/cpu"
priority = "explicit"

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