Second order optimizers

[Drenderer]

Optax doesn't support second-order optimisers such as SLSQP at the moment. We should investigate alternatives such as JAXopt or Optimistix and discuss potential integrations into klax.fit.

[jaosch]

JAXOpt appears to be deprecated, see here, but including other optimizer frameworks could certainly be an option.