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.