pyproject.toml

[tool.poetry]
name = "saex"
version = "0.0.1"
description = "SAEs in Jax"
authors = ["nev"]
license = "MIT"
readme = "README.md"

[tool.poetry.dependencies]
python = "^3.10"
numpy = "^1.26.4"
jax = {version = "^0.4.30", extras = ["tpu"]}
einops = "^0.7.0"
wandb = "^0.16.6"
matplotlib = "^3.8.4"
tqdm = "^4.66.2"
jaxtyping = "^0.2.28"
optax = "^0.2.2"
equinox = {git = "https://github.com/neverix/equinox.git", rev="42f324836138cd27630674aab82bc3042638065c"}
transformers = "^4.39.3"
datasets = "^2.18.0"
flax = "^0.8.2"
safetensors = "^0.4.3"
tokenizers = "^0.19.1"
torch = {version = "^2.2.2+cpu", source = "torch-cpu"}
jax-smi = {git = "https://github.com/neverix/jax-smi.git", rev="eb75686f31258a4cc05c1ddc7723d74e302a0ecd"}
oryx = "^0.2.6"
fire = "^0.6.0"
micrlhf-progress = {git = "https://github.com/neverix/micrlhf-progress.git"}
tensorflow = "^2.16.1"
tensorboard-plugin-profile = "^2.15.1"
more-itertools = "^10.2.0"
tensorflow-probability = {extras = ["tf"], version = "^0.24.0"}
jmp = "^0.0.4"
aqtp = "^0.7.5"

[tool.poetry.dev-dependencies]
pytest = "^7.1"
flake8 = "^4.0"

[[tool.poetry.source]]
name = "google-libtpu"
url = "https://storage.googleapis.com/jax-releases/libtpu_releases.html"
priority = "supplemental"


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

[tool.poetry.group.dev.dependencies]
isort = "^5.13.2"
ipykernel = "^6.29.4"
astor = "^0.8.1"
plotly = "^5.21.0"
ipympl = "^0.9.4"
nbformat = "^5.10.4"
line-profiler = "^4.1.3"
huggingface-hub = "^0.23.0"
gradio = "^4.31.5"

[tool.flake8]
max-line-length = 120

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