Use jax.value_and_grad for energy operator

16 jobs for jax_operator in 16 minutes and 18 seconds (queued for 1 second)