CIRQ How to iteratively apply a multi qubit gate to first n qubits

Question

I've got an arbitrary n qubit circuit, with a "for Q in range(n):", which creates a custom gate class that affects (Q+1) qubits, which I want to apply to first (Q+1) qubits of the circuit until the last gate that is applied to all qubits. I'm having trouble with the cirquit.append() command for this use case:

qubit = cirq.LineQubit.range(n_qubits)
circuit = cirq.Circuit()

for i in range(n_qubits):
    class CustomGate(cirq.Gate):
        def __init__(self):
            super(CustomGate, self)

        def _num_qubits_(self):
            return i+1

        def _unitary_(self):
            return np.array(all_gates[i])

        def _circuit_diagram_info_(self, args):
            return ["Gate" + str(i)] * (i+1)
    
    circuit.append(CustomGate().on(*qubit[:i+1]))  #HERE IS THE ERROR

It applies the first iteration to qubit[0] just fine, but at two it gives the following error:

ValueError: Wrong number of qubits for <<__main__.CustomGate object at 0x00000252D387F580>>. Expected 2 qubits but got <[cirq.LineQubit(0)]>.

Interestingly, I've tried to do a botched approach with if statements:

if i == 0:
        circuit.append(CustomGate().on(cirq.LineQubit(0)))
if i == 1:
        circuit.append(CustomGate().on(cirq.LineQubit(0), cirq.LineQubit(1)))
if i == 2:
        circuit.append(CustomGate().on(cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)))

But that breaks down at i==1 step when I try to print the circuit at that step.

But if I remove the for loop and write it out by hand three times it magically works and I've got no clue why:

class CustomGate(cirq.Gate):
        def __init__(self):
            super(CustomGate, self)

        def _num_qubits_(self):
            return 1

        def _unitary_(self):
            return np.array(all_gates[0])

        def _circuit_diagram_info_(self, args):
            return ["Gate" + str(0)] * (0+1)

circuit.append(CustomGate().on(cirq.LineQubit(0)))

class CustomGate(cirq.Gate):
        def __init__(self):
            super(CustomGate, self)

        def _num_qubits_(self):
            return 2

        def _unitary_(self):
            return np.array(all_gates[1])

        def _circuit_diagram_info_(self, args):
            return ["Gate" + str(1)] * (1+1)

circuit.append(CustomGate().on(cirq.LineQubit(0), cirq.LineQubit(1)))

class CustomGate(cirq.Gate):
        def __init__(self):
            super(CustomGate, self)

        def _num_qubits_(self):
            return 3

        def _unitary_(self):
            return np.array(all_gates[2])

        def _circuit_diagram_info_(self, args):
            return ["Gate" + str(2)] * (2+1)

circuit.append(CustomGate().on(cirq.LineQubit(0), cirq.LineQubit(1), cirq.LineQubit(2)))

None of these issues come up when the custom gate acts on one qubit instead of multiple. Please help, I've been stuck on this for hours

Answer 1

The problem is a python issue - you keep redefining the CustomGate class in the namespace, which can have weird consequences. Typically classes are defined once in a more global scope, not within loops.

You can see the bug if you print out the _num_qubits_() for each operation in the circuit after each iteration:

qubit = cirq.LineQubit.range(n_qubits)
circuit = cirq.Circuit()

for i in range(n_qubits):
    class CustomGate(cirq.Gate):
        def __init__(self):
            super(CustomGate, self)

        def _num_qubits_(self):
            return i+1

        def _unitary_(self):
            return np.array(all_gates[i])

        def _circuit_diagram_info_(self, args):
            return ["Gate" + str(i)] * (i+1)    
    circuit.append(CustomGate().on(*qubit[:i+1]))  
    print([op.gate._num_qubits_() for op in circuit.all_operations()])

Prints:

[1]
[2, 2]
[3, 3, 3]
[4, 4, 4, 4]
[5, 5, 5, 5, 5]

So, for example in the last iteration, all operations think that they have 5 qubits, but instead they have less!

Here's a fix that works:

n_qubits = 5
qubit = cirq.LineQubit.range(n_qubits)
circuit = cirq.Circuit()
class CustomGate(cirq.Gate):
      def __init__(self, i):
          super(CustomGate, self)
          self.i = i

      def _num_qubits_(self):
          return self.i+1

      def _unitary_(self):
          return np.array(all_gates[self.i])

      def _circuit_diagram_info_(self, args):
          return ["Gate" + str(self.i)] * (self.i+1)

for i in range(n_qubits):       
    circuit.append(CustomGate(i).on(*qubit[:i+1]))  

print(circuit)

Which prints:

0: ───Gate0───Gate1───Gate2───Gate3───Gate4───
              │       │       │       │
1: ───────────Gate1───Gate2───Gate3───Gate4───
                      │       │       │
2: ───────────────────Gate2───Gate3───Gate4───
                              │       │
3: ───────────────────────────Gate3───Gate4───
                                      │
4: ───────────────────────────────────Gate4───