numpy
flax
jax
numpyro
optax
matplotlib
pyswarms
scipy
tqdm
jaxtyping
seaborn
