jax>=0.7.0
jaxlib>=0.7.0
multipledispatch
numpy
tqdm

[cpu]
jax[cpu]>=0.7.0

[cuda12]
jax[cuda12]>=0.7.0

[cuda13]
jax[cuda13]>=0.7.0

[dev]
dm-haiku>=0.0.14
equinox
flax
funsor>=0.4.1
graphviz
jaxns<=2.6.9,>=2.6.3
matplotlib
optax>=0.0.6
pylab-sdk
pytest-cov
pyyaml
requests
tfp-nightly

[doc]
ipython
nbsphinx>=0.8.9
readthedocs-sphinx-search>=0.3.2
sphinx>=5
sphinx_rtd_theme
sphinx-gallery

[examples]
arviz
jupyter
matplotlib
pandas
seaborn
scikit-learn
wordcloud

[test]
importlib-metadata<5.0
ruff>=0.1.8
mypy>=1.13
pytest>=4.1
pyro-api>=0.1.1
scikit-learn
scipy>=1.9
ty>=0.0.4

[tpu]
jax[tpu]>=0.7.0
