jax>=0.2.16
flax
audax
torch
