Open KetpuntoG opened 2 months ago
Issue seems to be with taking parameter shift of a GlobalPhase
:
This gives the exact same error:
import numpy as np
import pennylane as qml
import jax
# with shots = None it works
dev = qml.device("default.qubit", wires = 4, shots=100)
@qml.qnode(dev)
def circuit(phase):
qml.GlobalPhase(phase)
return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
params = jax.numpy.array(0.4)
jac_fn = jax.jacobian(circuit)
# without jax.jit it works
jac_fn = jax.jit(jac_fn)
jac = jac_fn(params)
print(jac)
Potentially need a custom gradient recipe or something. Or just identify that the gradient of a global phase is always zero.
Hey @albi3ro and @KetpuntoG, I was looking at issues to contribute to. Can I tackle this issue ?
Just setting GlobalPhase.grad_method = None
seems to work 🤞 . That would basically just indicate "no differentiable parameters here". Should be a simple enough fix if you're interested.
In the operator I'm working on, I am differentiating with respect to the global phase parameter, would it stop working @albi3ro ?
Using a controlled global phase or something?
Correct 👌 It seems strange, but it appears naturally in the formulation
We could potentially update Controlled.grad_method
from:
@property
def grad_method(self):
return self.base.grad_method
to:
@property
def grad_method(self):
return "A" if self.base.name == "GlobalPhase" else self.base.grad_method
Hey @albi3ro and @KetpuntoG , I want to work on this. Can you please assign it to me.
Assigned. Note that our next release is in two weeks on May 6th, so we may take it over next week to make sure we get the fix in.
Hey @albi3ro , I am unsure how to add tests for this. The below test fails as the values are far apart
@pytest.mark.jax
def test_jacobian_with_and_without_jit_has_same_output():
import jax
dev = qml.device("default.qubit", wires=4, shots=100)
@qml.qnode(dev)
def circuit(coeffs):
qml.AmplitudeEmbedding(coeffs, normalize=True, wires=[0, 1])
return qml.expval(qml.PauliZ(0) @ qml.PauliZ(1))
params = jax.numpy.array([0.4, 0.5, 0.1, 0.3])
jac_fn = jax.jacobian(circuit)
jac_jit_fn = jax.jit(jac_fn)
# Array([ 1.67562983, -1.88046271, -0.48002058, 1.05993827], dtype=float64
jac = jac_fn(params)
# Array([ 1.77878858, -2.02651428, -0.11825829, 1.04522512], dtype=float64))
jac_jit = jac_jit_fn(params)
assert qml.math.allclose(jac, jac_jit)
I just made GlobalPhase.grad_method = None
as mentioned in this comment
Sorry about not getting back to you earlier @Tarun-Kumar07 . Been a bit busy the last few days.
This looks to be a case of the shots being too low. I bumped it up to shots=5000
, and then the numbers started to converge better.
So there's two options:
1) Just bumping up the shot number and setting a seed on the device to reduce flakiness qml.device('default.qubit', shots=10000, seed=7890234)
2) Using analytic mode shots=None
and manually specifying diff_method="parameter-shift"
.
Potentially we can just test both.
Hey @albi3ro ,
I tried writing tests for MottonenStatePreparation
and StatePrep
similar to the one in the above comment. Even by providing a high number of shots and using the analytic mode, these tests fail.
I have opened a draft PR: #5620, and these are the only tests failing, not sure how to fix them.
Expected behavior
qml.AmplitudeEmbedding
should work withjit
,grad
and finite shotsActual behavior
ValueError: need at least one array to stack
Additional information
Same issue with
qml.StatePrep
andqml.MottonenStatePreparation
Source code
No response
Tracebacks
No response
System information
Existing GitHub issues