flashoptim_jax
