Draw a qnode subroutine as a block

I am currently implementing Shor’s algorithm with a custom modular multiplication matrix.
It doesn’t work well, so I want to draw it to have an idea of what the problem could be. However, the circuit is way too big to be fully displayed, because my multiplications use several additions, some cleanup routines, etc…

So I was wondering something: is it possible to draw a subroutine (like, let’s say, an addition circuit) as a block in the global circuit drawing ? That would be like for example the QFT, which is a complex circuit but in Pennylane drawing we just see a block “QFT”.

Thanks.

Hi @Asimonu, welcome to the forum!

I think the easiest way would be to define a custom gate. In this blog post you can find the details on how to make one.

Here is a working example:

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

class RXX(Operation):
    num_params = 1
    num_wires = 2
    par_domain = "R"

    @staticmethod
    def decomposition(theta, wires):
        return [qml.PauliRot(theta, 'XX', wires=wires)]
    
    dev = qml.device('default.qubit', wires=2, shots=100)

@qml.qnode(dev)
def circuit(theta):
    RXX(theta, wires=[0, 1])
    qml.Hadamard(1)
    return qml.expval(qml.PauliZ(0) @ qml.PauliY(1))

qml.draw_mpl(circuit)(0.2)

Please let me know if this works for you!

1 Like

Yes, this solution works very well, thank you!

I’m glad it works @Asimonu.

Enjoy using PennyLane!