numba
dynamax>=0.1.4
chex
jaxtyping
tensorflow_probability
tqdm
pandas
ipykernel
jupyter
h5py
matplotlib
vidio
imageio[ffmpeg]
scikit-learn
plotly

[cuda]
jax[cuda12]

[dev]
sphinx
sphinx-rtd-theme
autodocsumm
myst-nb
