jax>=0.7.0
optax>=0.1.7
orbax-checkpoint>=0.4.0
flax>=0.7.5
chex>=0.1.84
jaxtyping>=0.2.23
numpy>=1.24.0
tqdm>=4.65.0

[all]
fabricpc[experiments,tfds,viz]

[cpu]
jax[cpu]

[cuda12]
jax[cuda12]

[cuda13]
jax[cuda13]

[dev]
pytest>=7.0.0
hypothesis>=6.0.0
black[colorama]==26.1.0
ruff==0.15.19
mypy>=1.0.0
pre-commit>=3.0.0
build>=1.0.0
twine>=5.0.0

[experiments]
scipy>=1.10.0
optuna>=3.0.0

[tfds]
tensorflow-datasets>=4.9.0
importlib_resources
tokenizers>=0.15.0

[tfds:platform_system != "Linux" or platform_machine != "x86_64"]
tensorflow>=2.15.0

[tfds:platform_system == "Linux" and platform_machine == "x86_64"]
tensorflow-cpu>=2.15.0

[viz]
plotly>=5.0.0
kaleido>=0.2.1
pandas>=2.0.0

[viz:python_version < "3.13" and platform_system != "Windows"]
aim>=3.0.0
