numpy
tqdm
torch
jax
numpyro
optax
