`parametrize_differentiable` not compatible with `custom_decomps`

[yitchen-tim]

Describe the bug pennylane circuits are un-parametrized if there the qnode has custom_decomps keyword argument.

To reproduce

import pennylane.numpy as np
import pennylane as qml

custom_decomps={"IsingXX": qml.IsingXX.compute_decomposition}

dev_local = qml.device(
    "braket.local.qubit",
    wires=2, 
    shots=1000,
    custom_decomps=custom_decomps,
    parametrize_differentiable=True
)

@qml.qnode(dev_local)
def my_pl_circuit(a, b):
    qml.RX(a, wires=0)
    qml.IsingXX(b, wires=[0,1])
    return [qml.expval(qml.PauliZ(i)) for i in range(2)]

a = np.array(0.1)
b = np.array(0.2)
_ = my_pl_circuit(a, b)

print(my_pl_circuit.device.circuit)

Expected behavior The underlying Braket circuit is a circuit parametrized by free parameters a and b.

Screenshots or logs The circuit is not parametrized. Even though custom_decomps only applies on the XX gate, the RX gate also loses parametrization and has the value of parameter a bound to it.

T  : |   0    |1|   2    |3| Result Types |

q0 : -Rx(0.10)-C-Rx(0.20)-C-Expectation(Z)-
               |          |                
q1 : ----------X----------X-Expectation(Z)-

T  : |   0    |1|   2    |3| Result Types |

If I remove the custom_decomps kwarg from qnode, the circuit become parametrized again.

System information A description of your system. Please provide:

• Amazon Braket Python PennyLane Plugin version: 1.19.2.dev0
• Amazon Braket Python SDK version: 1.54.0
• Amazon Braket Python Schemas version: 1.19.1
• Amazon Braket Python Default Simulator version: 1.20.0
• Python version: 3.9.7

Additional context Add any other context about the problem here.

[speller26]

Turns out that when you supply the custom_decomps argument to the device, tape.get_parameters() and tape.trainable_params become empty; interesting behavior from PennyLane.

[albi3ro]

So I can't find the correct branch to try this out with, but I can recommend a temporary user-fix:

@qml.qnode(dev_local, expansion_strategy="device")
def my_pl_circuit(a, b):
    qml.RX(a, wires=0)
    qml.IsingXX(b, wires=[0,1])
    return [qml.expval(qml.PauliZ(i)) for i in range(2)]

This way the device expansion will happen earlier on when the parameters are still autograd.

The device expansion is probably wiping out any information we have about trainability. Normally this doesn't matter. If a gradient transform is using the trainability information, then the conversion to numpy and the device expansion happens after teh gradient transform is done with it. If a device gradient is using the trainability information, then device expansion happens before we hit the ml framework boundary, set the correct trainable parameters, and convert to numpy.

We can look into better fixes, but I'm not sure they would be particularly easy to figure out.

[speller26]

Thanks for the suggesion @albi3ro! As for the correct branch to try this out, you can just use main!

[albi3ro]

I just ran the above code locally and got:

T  : |   0   |1|   2   |3| Result Types |

q0 : -Rx(p_0)-C-Rx(p_1)-C-Expectation(Z)-
              |         |                
q1 : ---------X---------X-Expectation(Z)-

T  : |   0   |1|   2   |3| Result Types |

Unassigned parameters: [p_0, p_1].

What should I be seeing once this is fixed?

[yitchen-tim]

Hi @albi3ro , the one you have is the expected output. Looks like it's fixed with Pennylane v0.33. We can close this issue, if @speller26 you don't see other issue.

[speller26]

Perfect, looks like the problem's solved! Thanks @albi3ro!