jax[cuda]
flax
jaxtyping
jaxlib
