Using torch.autograd on circuit containing Kerr gate on strawberryfields.fock

PennyLane Help
1

Dear Pennylane Team,

I have a few questions regarding the use of torch.autograd on a continuous variable circuit runnnig on the strawberryfields.fock device.

  1. Am I correct that it is not possible to use torch.autograd.grad on a (quantum) torch.sequential as long as the circuit contains a non-linear operation (as the Kerr gate)?

The following small example does not work when the circuit includes a Kerr gate but perfectly fine otherwise:

from pennylane import expval, X
from pennylane.qnn import TorchLayer
import pennylane as qml
import numpy as np
import torch

nqubits = 2
cutoff_dim = 4

device = qml.device('strawberryfields.fock', wires=nqubits, 
cutoff_dim=cutoff_dim)

@qml.qnode(device=device, interface='torch', diff_method="best")
def cvqcFkt(inputs, w0, w1):
    qml.Displacement(inputs[0], 0, wires=0)
    qml.Rotation(w0, wires=0)
    qml.Rotation(w1, wires=0)
    #### Comment out Kerr gate t make program work ####
    qml.Kerr(w0, wires=0)
    qml.Squeezing(w1, 0.0, wires=0)
    return [expval(X(wires=wire)) for wire in range(nqubits)]

weight_shapes = {"w0": (1), "w1": (1)}

qlayer = qml.qnn.TorchLayer(cvqcFkt, weight_shapes)
qlayers = torch.nn.Sequential(qlayer)

x = torch.tensor([0.1], dtype=float, requires_grad=True)
u = qlayers(x)

u_x = torch.autograd.grad(
     u, x,
     grad_outputs=torch.ones_like(u, requires_grad=True),
     retain_graph=True,
    create_graph=True,
    allow_unused=True
    )[0]

print(u_x)

The error I get (when there is a KerrGate) is :

TypeError: all() received an invalid combination of arguments - got 
(axis=NoneType, out=NoneType, ), but expected one of:
 * ()
 * (int dim, bool keepdim)
      didn't match because some of the keywords were incorrect: axis, 
out
 * (name dim, bool keepdim)
      didn't match because some of the keywords were incorrect: axis, 
out

I saw that differentiating CV circuits with Kerr gates is indeed possible with tensorflow but as most of our program relies on pytorch, I would love to find a solution within this framework. So my further questions:

  1. Are there workarounds in pytorch?
  2. Is this planned to run in the (closer) future?

If you have questions regarding by my questions, please let me know.

Thanks a lot! and best regards,
Pia

2

Hi @Pia!

Thank you for reporting this, this looks to be a bug in PennyLane when using PyTorch; I’ve made a bugfix here: https://github.com/PennyLaneAI/pennylane/pull/2117.

In the meantime, you could set the max_diff argument to 2 when creating your QNode; this should avoid the bug:

@qml.qnode(device, interface='torch', max_diff=2)

One final note: unfortunately, the parameter-shift rule (which returns exact gradients) does not support non-linear gates such as the Kerr gate. So to compute the derivative of the Kerr gate parameter (and any gates that follow), PennyLane will fallback to finite-differences.

Switching to "strawberryfields.tf" and interface="tf" does have an advantage in this case — PennyLane will be able to use backpropgation for exact gradients, as the underlying device uses TensorFlow.

1 Like
3

Hi @josh!

Thanks a lot for the fast and helpful reply. Setting max_diff=2 does indeed make it work.

Best regards,
Pia