catalyst.draw_graph

draw_graph(qnode: QJIT, *, level: int | None = None) → Callable

Visualize a single QJIT compiled QNode, showing wire flow through quantum operations, program structure, and pass-by-pass impacts on compiled programs.

Warning

This function only visualizes quantum operations contained in workflows involving a single qjit-compiled QNode. Workflows involving multiple QNodes or operations outside QNodes cannot yet be visualized.

Only transformations found within the Catalyst compiler can be visualized. Any PennyLane tape transform will have already been applied before lowering to MLIR and will appear as the base state (level=0) in this visualization.

Note

The draw_graph function visualizes a QJIT-compiled QNode in a similar manner as view-op-graph does in MLIR, which leverages Graphviz to show data-flow in the compiled IR.

As such, use of draw_graph requires installation of Graphviz and the pydot software package. Please consult the links provided for installation instructions.

Additionally, it is recommended to use draw_graph with PennyLane’s program capture enabled (see qml.capture.enable).

Parameters:

• qnode (QJIT) – The input qjit-compiled QNode that is to be visualized. The QNode is assumed to be compiled with qjit.

• level (int | None) – The level of transformation to visualize. If None, the final level is visualized.

Returns:

A function that has the same argument signature as the compiled QNode. When called, the function will return the graph as a tuple of (matplotlib.figure.Figure, matplotlib.axes._axes.Axes) pairs.

Return type:

Callable

Raises:

VisualizationError – If the circuit contains operations that cannot be converted to a graphical representation.

Example

Using draw_graph requires a qjit’d QNode and a level argument, which denotes the cumulative set of applied compilation transforms (in the order they appear) to be applied and visualized.

import pennylane as qml
import catalyst

qml.capture.enable()

@qml.qjit
@qml.transforms.merge_rotations
@qml.transforms.cancel_inverses
@qml.qnode(qml.device("null.qubit", wires=3))
def circuit():
    qml.H(0)
    qml.T(1)
    qml.H(0)
    qml.RX(0.1, wires=0)
    qml.RX(0.2, wires=0)
    return qml.expval(qml.X(0))

With level=0, the graphical visualization will display the program as if no transforms are applied:

>>> print(catalyst.draw_graph(circuit, level=0)())
(<Figure size 640x480 with 1 Axes>, <Axes: >)

With level=2, both merge_rotations() and cancel_inverses() will be applied, resulting in the two Hadamards cancelling and the two rotations merging:

>>> print(catalyst.draw_graph(circuit, level=2)())
(<Figure size 640x480 with 1 Axes>, <Axes: >)

Usage Details

Visualizing Control Flow

The draw_graph function can be used to visualize control flow, resulting in a scalable representation that preserves program structure:

@qml.qjit(autograph=True)
@qml.qnode(qml.device("null.qubit", wires=3))
def circuit():
    qml.H(0)
    for i in range(3):
        if i == 1:
            qml.X(0)
        elif i == 2:
            qml.Y(0)
        else:
            qml.Z(0)
    return qml.probs()

>>> print(catalyst.draw_graph(circuit)())
(<Figure size 640x480 with 1 Axes>, <Axes: >)

As one can see, the program structure is preserved in the figure.

Visualizing Dynamic Circuits

Circuits can depend on parameters that are not known at compile time, which result in conventional visualization tools failing. Consider the following circuit.

@qml.qjit
@qml.qnode(qml.device("null.qubit", wires=3))
def circuit(x, y):
    qml.X(0)
    qml.Y(1)
    qml.Z(2)
    qml.H(x) # 'x' is a dynamic wire index
    qml.S(0)
    qml.T(2)
    qml.H(x)
    return qml.expval(qml.Z(y))

The two qml.H gates act on wires that are dynamic. In order to preserve qubit data flow, each dynamic operator acts as a “choke point” to all currently active wires. To visualize this clearly, we use dashed lines to represent a dynamic dependency and solid lines for static/known values:

>>> x, y = 1, 0
>>> print(catalyst.draw_graph(circuit)(x, y))
(<Figure size 640x480 with 1 Axes>, <Axes: >)

code/api/catalyst.draw_graph

 

Download Python script

 

Download Notebook

 

View on GitHub