jaxtyping>=0.2.19
jax>=0.4.0
flax>=0.7.0
pgx>=2.0.1
tqdm~=4.66.0
