jax
equinox
tqdm
