jax
jaxlib
jaxtyping
optax
beartype

[all]
equinox
flax
datasets
tqdm
pillow
wandb

[all:python_version >= "3.14"]
please-downgrade-to-python-3.13-for-flax

[equinox]
equinox

[flax]
flax

[flax:python_version >= "3.14"]
please-downgrade-to-python-3.13-for-flax

[test]
pytest

[train]
datasets
tqdm
pillow
wandb
