How can I implement a repeated (non-unitary) custom gate with mid-circuit measurements?

I’m trying to implement a non-unitary operation, which is made deterministic by post-selection on a mid-circuit measurement. This works as expected using the following code:

import pennylane as qml
from pennylane import numpy as np
from pennylane.operation import Operation, AnyWires

from numpy import tan, tanh, arctan, arctanh, pi, cosh, sinh

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

@qml.qnode(dev)
def single_layer_circuit(tau):
    t = 2*arctan(tanh(tau))
    qml.Hadamard(wires=0)
    qml.Hadamard(wires=1)
    qml.MultiRZ(t, wires=[0,1])
    qml.Hadamard(wires=0)
    qml.RX(-pi/2, wires=0)
    qml.measure(0, postselect=0)
    return qml.density_matrix(1)

Now I wish to repeat the application of this gate in a compact manner. I’m trying to do this by using qml.layer, as follows:

import pennylane as qml
from pennylane import numpy as np
from pennylane.operation import Operation, AnyWires

from numpy import tan, tanh, arctan, arctanh, pi, cosh, sinh

def single_layer(tau):
    t = 2*arctan(tanh(tau))
    qml.Hadamard(wires=0)
    qml.Hadamard(wires=1)
    qml.MultiRZ(t, wires=[0,1])
    qml.Hadamard(wires=0)
    qml.RX(-pi/2, wires=0)
    qml.measure(0, postselect=0)

@qml.qnode(dev)
def circuit(params):
    qml.layer(single_layer, 2, params)
    return qml.density_matrix(1)

params = np.array([[0.5], [0.4]])

This works if I remove the mid-circuit measurement, meaning that the most straightforward way to repeat this gate (with possibly different parameters) requires some additional loop, making the code much more cumbersome.

Is there anyway to make this gate (with measurements) iterative via qml.layer? Or in some other way which makes it easy to manage the different parameters in the layers (e.g. by defining some Operation class)? Thank you.

Hey @physics_student, welcome to the forum! :sunglasses:

The issue is the way you’ve shaped params and how that interferes with MultiRZ. If params has no leading broadcasting dimension (i.e., it’s 1D), this works:

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit")

def single_layer(params):
    qml.MultiRZ(params, wires=[0, 1])
    qml.measure(0, postselect=0)

@qml.qnode(dev)
def circuit(params):
    qml.layer(single_layer, 2, params)
    return qml.density_matrix(1)

params = np.array([0.5, 0.4])

circuit(params)
tensor([[1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j]], requires_grad=True)

However, using this will throw an error:

params = np.array([[0.5], [0.4]]) # shape is (2, 1) — has a leading dimension
circuit(params)
ValueError: Cannot postselect on circuits with broadcasting. Use the qml.transforms.broadcast_expand transform to split a broadcasted tape into multiple non-broadcasted tapes before executing if postselection is used.

However, this is strange; if I replace MultiRZ with RZ and try to parameter broadcast, it works??

import pennylane as qml
from pennylane import numpy as np

dev = qml.device("default.qubit")

def single_layer(params):
    #qml.MultiRZ(params, wires=[0, 1])
    qml.RZ(params, wires=[0])
    qml.Hadamard(1)
    qml.measure(0, postselect=0)

@qml.qnode(dev)
def circuit(params):
    qml.layer(single_layer, 2, params)
    return qml.density_matrix(1)

params = np.array([0.5, 0.4])

circuit(params)
tensor([[1.+0.j, 0.+0.j],
        [0.+0.j, 0.+0.j]], requires_grad=True)

I think for now, your issue will be solved, but you might have uncovered a bug accidentally? I’ll check with the team!

1 Like

Hey @physics_student,

Glad you found my last post helpful! Just wanted to follow-up. There are no bugs, but here are some examples that may illustrate why there’s an error:

def single_layer(params):
    qml.MultiRZ(params, wires=[0, 1])
    qml.measure(0, postselect=0)

@qml.qnode(dev)
def circuit(params):
    qml.layer(single_layer, 2, params)
    return qml.density_matrix(1)

params = np.array([0.5, 0.4])

print(circuit(params))
print(qml.draw(circuit)(params))
[[1.+0.j 0.+0.j]
 [0.+0.j 0.+0.j]]
0: ─╭MultiRZ(0.50)──┤↗₀├─╭MultiRZ(0.40)──┤↗₀├─┤       
1: ─╰MultiRZ(0.50)───────╰MultiRZ(0.40)───────┤  State

This is the expected output. If you create a qml.layer where the layer is repeated twice, the params you give it must be length 2, where each entry is the argument provided to the single layer that gets repeated.

def single_layer(params):
    qml.MultiRZ(params, wires=[0, 1])
    qml.measure(0, postselect=0)

@qml.qnode(dev)
def circuit(params):
    qml.layer(single_layer, 2, params)
    return qml.density_matrix(1)

params = np.array([[0.5], [0.4]]) # <<<<<< This is different!

print(circuit(params))
print(qml.draw(circuit)(params))
ValueError: Cannot postselect on circuits with broadcasting. Use the qml.transforms.broadcast_expand transform to split a broadcasted tape into multiple non-broadcasted tapes before executing if postselection is used.

You’ve now introduced a leading dimension in params! It has a shape of (2, 1), where 2 is the leading dimension. PennyLane sees this and wants to broadcast because MultiRZ only needs one argument. You can get around this with broadcast_expand, but I’m not sure this is what you intended to do in the first place :). Just be careful with leading dimensions!