[BUG] QubitStateVector is decomposed even on devices that support it

Expected behavior

QubitStateVector should not be decomposed and the output should be:

 0: ──╭QubitStateVector(M0)──╭┤ State 
 1: ──╰QubitStateVector(M0)──╰┤ State 
M0 =
[1. 0. 0. 0.]

This happens up to 0.19.1

Actual behavior

QubitStateVector is decomposed

 0: ──╭C──╭C──╭┤ State 
 1: ──╰X──╰X──╰┤ State 

This happens in 0.20.0 an beyond.

Additional information

No response

Source code

import pennylane as qml
from pennylane.devices import DefaultQubit
from pennylane import numpy as np

dev = qml.device('default.qubit', wires=2)

@qml.qnode(dev, diff_method='parameter-shift')
def circuit():
    qml.QubitStateVector(np.array([1.0, 0.0, 0.0, 0.0]), wires=[0, 1])
    return qml.state()

assert 'QubitStateVector' in dev.operations
print(qml.draw(circuit)())

Tracebacks

No response

System information

As stated above the problem starts with 0.20.0.

Existing GitHub issues

• I have searched existing GitHub issues to make sure the issue does not already exist.

[josh146]

Yeah, that was my idea at the time. It ended up failing because if you call the gradient decompositions inside the gradient logic, you can enter a case where classical processing occurs, which can’t be tracked by the autodiff framework; so you end up with the incorrect gradients being returned.

E.g., if you have a decomposition that looks like

Gate1(x) -> [Gate2(x^2), Gate3(sin(x))]

— that is, a quantum gate is decomposed down to quantum + classical components — the parameter-shift rule would only give the quantum component of the gradient, the classical component is lost.

I tried several workarounds, one being manually tracking the classical component of the decomposition, and returning class_jac @ quantum_jac, but this ended up adding a lot of complexity and overhead 🤔 I also couldn’t get it to work for all edge cases at the time, if I recall.

Moving the gradient-based decomposition higher up the stack, where the autodiff frameworks can continue to track classical processing, solves this, with the disadvantage being as we discussed above.

But this is definitely something we would love to support; being able to compute hybrid Jacobians without relying on the autodiff frameworks would solve this.

[cvjjm]

I see. But couldn’t the decomposition be done in the jacobian method of the respective tapes in a way that works for each interface?