Thanks for this @Tarun-Kumar07

For the failures due to errors:

ValueError: Computing the gradient of circuits that return the state with the parameter-shift rule gradient transform is not supported, as it is a hardware-compatible method.

That would be expected, and we should shift the measurement to expectation values.

For the failures due to:

FAILED tests/templates/test_state_preparations/test_mottonen_state_prep.py::test_jacobian_with_and_without_jit_has_same_output_with_high_shots[StatePrep] - AssertionError: assert Array(False, dtype=bool)
 +  where Array(False, dtype=bool) = <function allclose at 0x7fdb77481940>(Array([-0.0003,  0.0153, -0.0153,  0.0003], dtype=float64), Array([ 1.0187, -0.9953, -1.0047,  0.9813], dtype=float64), atol=0.02)

Those are legitimately different results, so we can safely safe we are getting wrong results in that case 😢 I'll investigate.

Those are legitimately different results, so we can safely safe we are getting wrong results in that case 😢 I'll investigate.

@Tarun-Kumar07 @albi3ro Not sure you got to this yet, but it seems that the decomposition of those state preparation methods handle special parameter values differently than others. This makes the derivative wrong at those special values, because param_shift is handed a tape that does not contain the general decomposition, so it will not shift all operations that need shifting. As JITting does not allow such special cases, we only make this mistake without JITting, hence the difference in the results within those tests.

Basically, the decomposition does something like the following decomposition for RZ:

def compute_decomposition(theta, wires):
    if not qml.math.is_abstract(theta) and qml.math.isclose(theta, 0):
        return []
    return [qml.RZ(theta, wires)]

It's correct but it does not have the correct parameter-shift derivative at 0.

This looks like an independent bug to me, and like one that could be hiding across the codebase for other ops as well, theoretically.