jax
jaxopt

[cuda]
jax[cuda12]
