jax>=0.2.16
flax
torch
audax
