jaxopt.GradientDescent() is not converging

it is my first encounter w/ PennyLane, so I’m probably making some trivial mistake.
I’m exercising this tutorial:

All works as expected, except the last example with optimization of the strength of the noise channel

gd = jaxopt.GradientDescent(cost, maxiter=steps)
res = gd.run(x, ev)

Is not converging. I see this answer:

QNode output after optimization = 0.5000
Experimental expectation value = 0.7781
Optimized noise parameter p = 0.5025

instead of

QNode output after optimization = 0.7771

Please advice,
Below is my software stack

$ pip3 list |grep jax
jax                       0.4.23
jaxlib                    0.4.23
jaxopt                    0.8.3

>>> print(qml.about())
Name: PennyLane
Version: 0.34.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/PennyLaneAI/pennylane


License: Apache License 2.0

Location: /usr/local/lib/python3.10/dist-packages

Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions

Required-by: PennyLane-Cirq, PennyLane-Lightning, PennyLane-Lightning-GPU, PennyLane-qiskit, PennyLane-SF

Platform info: Linux-6.6.8-200.fc39.x86_64-x86_64-with-glibc2.35

Python version: 3.10.12

Numpy version: 1.26.3

Scipy version: 1.11.4

Installed devices:

- default.gaussian (PennyLane-0.34.0)

- default.mixed (PennyLane-0.34.0)

- default.qubit (PennyLane-0.34.0)

- default.qubit.autograd (PennyLane-0.34.0)

- default.qubit.jax (PennyLane-0.34.0)

- default.qubit.legacy (PennyLane-0.34.0)

- default.qubit.tf (PennyLane-0.34.0)

- default.qubit.torch (PennyLane-0.34.0)

- default.qutrit (PennyLane-0.34.0)

- null.qubit (PennyLane-0.34.0)

- cirq.mixedsimulator (PennyLane-Cirq-0.34.0)

- cirq.pasqal (PennyLane-Cirq-0.34.0)

- cirq.qsim (PennyLane-Cirq-0.34.0)

- cirq.qsimh (PennyLane-Cirq-0.34.0)

- cirq.simulator (PennyLane-Cirq-0.34.0)

- lightning.qubit (PennyLane-Lightning-0.34.0)

- lightning.gpu (PennyLane-Lightning-GPU-0.34.0)

- strawberryfields.fock (PennyLane-SF-0.29.0)

- strawberryfields.gaussian (PennyLane-SF-0.29.0)

- strawberryfields.gbs (PennyLane-SF-0.29.0)

- strawberryfields.remote (PennyLane-SF-0.29.0)

- strawberryfields.tf (PennyLane-SF-0.29.0)

- qiskit.aer (PennyLane-qiskit-0.34.0)

- qiskit.basicaer (PennyLane-qiskit-0.34.0)

- qiskit.ibmq (PennyLane-qiskit-0.34.0)

- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.34.0)

- qiskit.ibmq.sampler (PennyLane-qiskit-0.34.0)

- qiskit.remote (PennyLane-qiskit-0.34.0)



I found a fix, copied from another tutorial:

This converges after ~150 iterations
The difference is in the initialization of the minimizer and its setting - I think.

opt = jaxopt.GradientDescent(cost, stepsize=0.4, acceleration = False)
steps = 200
params = x
opt_state = opt.init_state(params)
for i in range(steps):
    params, opt_state = opt.update(params, opt_state,target=target_ev)
    if (i + 1) % 10 == 0:
        print("Cost after step {:5d}: {: .7f}".format(i + 1, cost(params,target=target_ev)))

print("Optimized x: {}".format(params))
print(f"Optimized noise parameter p = {sigmoid(params.take(0)):.4f}")
print(f"QNode output after optimization = {circuit(params):.4f}")
print(f"Experimental expectation value = {target_ev}")

Hey @Jan_Balewski, welcome to the forum! Apologies for the delayed response.

I just want to make sure that the Noisy circuit demo is still functioning as it should. If you just run the code in the demo without changing anything, do you still get different results from what’s presented in the demo?

The answer is negative - the demo is not returning correct value ‘after optimization’.
To be sure I downloaded again the code as plain python using
and executed it on my laptop. Below is the output:

core@5b5c55ccbef0:~/toys$ python3 tutorial_noisy_circuits.py 
QNode output = 1.0000
Output density matrix is = 
[[0.5 0.  0.  0.5]
 [0.  0.  0.  0. ]
 [0.  0.  0.  0. ]
 [0.5 0.  0.  0.5]]
QNode output for bit flip probability 0.001 is 0.9960
QNode output for bit flip probability 0.01 is 0.9604
QNode output for bit flip probability 0.1 is 0.6400
QNode output for bit flip probability 0.2 is 0.3600
Output state for bit flip probability 0.2 is 
[[0.34 0.   0.   0.34]
 [0.   0.16 0.16 0.  ]
 [0.   0.16 0.16 0.  ]
 [0.34 0.   0.   0.34]]
QNode output for depolarizing probability 0.001 is 0.9973
QNode output for depolarizing probability 0.01 is 0.9735
QNode output for depolarizing probability 0.1 is 0.7511
QNode output for depolarizing probability 0.2 is 0.5378
QNode output after optimization = 0.5000  <== NOT CORRECT 
Experimental expectation value = 0.7781
Optimized noise parameter p = 0.5025

FYI, this is the Dockerfile constructing the software stack I’m using for this exercise:
Thanks, Jan

Hello, thank you for opening this question! On the website I am seeing that it converges to the correct solution. Can you confirm that the value of 0.5 you were seeing on the web?

I’m not sure I understand the question. Yes, on the website I also see value of 0.771, but it is just a static display. If I run the code I see the value of 0.5 - on my laptop.
So I can’t reproduce this tutorial as-is.

@Jan_Balewski this is strange :thinking: I ran this in a google collab notebook and was able to replicate what you see.

I will flag this to our team and get back to you :slight_smile:

Oh I see! Thanks for the clarification
In google collab works if you use jaxopt==0.5.5 that is the version we use in the website. In any case, I will communicate this behavior so that the team is aware of it :ok_hand: