How to optimize a QML model using JAX and JAXopt