From 265cd4d56ce067832966cd77fe807cd658ce6da1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 13:28:34 -0500 Subject: [PATCH 001/245] copy files over --- .../visualization/dag_builder.py | 100 ++++++ .../visualization/pydot_dag_builder.py | 207 ++++++++++++ .../visualization/test_dag_builder.py | 94 ++++++ .../visualization/test_pydot_dag_builder.py | 300 ++++++++++++++++++ 4 files changed, 701 insertions(+) create mode 100644 frontend/catalyst/python_interface/visualization/dag_builder.py create mode 100644 frontend/catalyst/python_interface/visualization/pydot_dag_builder.py create mode 100644 frontend/test/pytest/python_interface/visualization/test_dag_builder.py create mode 100644 frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py new file mode 100644 index 0000000000..7a60a55f1e --- /dev/null +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -0,0 +1,100 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""File that defines the DAGBuilder abstract base class.""" + +from abc import ABC, abstractmethod +from typing import Any + + +class DAGBuilder(ABC): + """An abstract base class for building Directed Acyclic Graphs (DAGs). + + This class provides a simple interface with three core methods (`add_node`, `add_edge` and `add_cluster`). + You can override these methods to implement any backend, like `pydot` or `graphviz` or even `matplotlib`. + + Outputting your graph can be done by overriding `to_file` and `to_string`. + """ + + @abstractmethod + def add_node( + self, node_id: str, node_label: str, parent_graph_id: str | None = None, **node_attrs: Any + ) -> None: + """Add a single node to the graph. + + Args: + node_id (str): Unique node ID to identify this node. + node_label (str): The text to display on the node when rendered. + parent_graph_id (str | None): Optional ID of the cluster this node belongs to. + **node_attrs (Any): Any additional styling keyword arguments. + + """ + raise NotImplementedError + + @abstractmethod + def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + """Add a single directed edge between nodes in the graph. + + Args: + from_node_id (str): The unique ID of the source node. + to_node_id (str): The unique ID of the destination node. + **edge_attrs (Any): Any additional styling keyword arguments. + + """ + raise NotImplementedError + + @abstractmethod + def add_cluster( + self, + cluster_id: str, + cluster_label: str, + parent_graph_id: str | None = None, + **cluster_attrs: Any, + ) -> None: + """Add a single cluster to the graph. + + A cluster is a specific type of subgraph where the nodes and edges contained + within it are visually and logically grouped. + + Args: + cluster_id (str): Unique cluster ID to identify this cluster. + cluster_label (str): The text to display on the cluster when rendered. + parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. + **cluster_attrs (Any): Any additional styling keyword arguments. + + """ + raise NotImplementedError + + @abstractmethod + def to_file(self, output_filename: str) -> None: + """Save the graph to a file. + + The implementation should ideally infer the output format + (e.g., 'png', 'svg') from this filename's extension. + + Args: + output_filename (str): Desired filename for the graph. + + """ + raise NotImplementedError + + @abstractmethod + def to_string(self) -> str: + """Return the graph as a string. + + This is typically used to get the graph's representation in a standard string format like DOT. + + Returns: + str: A string representation of the graph. + """ + raise NotImplementedError diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py new file mode 100644 index 0000000000..ad9eef4bcb --- /dev/null +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -0,0 +1,207 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""File that defines the PyDotDAGBuilder subclass of DAGBuilder.""" + +import pathlib +from collections import ChainMap +from typing import Any + +from .dag_builder import DAGBuilder + +has_pydot = True +try: + import pydot +except ImportError: + has_pydot = False + + +class PyDotDAGBuilder(DAGBuilder): + """A Directed Acyclic Graph builder for the PyDot backend.""" + + def __init__( + self, + attrs: dict | None = None, + node_attrs: dict | None = None, + edge_attrs: dict | None = None, + cluster_attrs: dict | None = None, + ) -> None: + """Initialize PyDotDAGBuilder instance. + + Args: + attrs (dict | None): User default attributes to be used for all elements (nodes, edges, clusters) in the graph. + node_attrs (dict | None): User default attributes for a node. + edge_attrs (dict | None): User default attributes for an edge. + cluster_attrs (dict | None): User default attributes for a cluster. + + """ + # Initialize the pydot graph: + # - graph_type="digraph": Create a directed graph (edges have arrows). + # - rankdir="TB": Set layout direction from Top to Bottom. + # - compound="true": Allow edges to connect directly to clusters/subgraphs. + # - strict=True": Prevent duplicate edges (e.g., A -> B added twice). + self.graph: pydot.Dot = pydot.Dot( + graph_type="digraph", rankdir="TB", compound="true", strict=True + ) + # Create cache for easy look-up + self._subgraphs: dict[str, pydot.Graph] = {} + self._subgraphs["__base__"] = self.graph + + _default_attrs: dict = {"fontname": "Helvetica", "penwidth": 2} if attrs is None else attrs + self._default_node_attrs: dict = ( + { + **_default_attrs, + "shape": "ellipse", + "style": "filled", + "fillcolor": "lightblue", + "color": "lightblue4", + "penwidth": 3, + } + if node_attrs is None + else node_attrs + ) + self._default_edge_attrs: dict = ( + { + "color": "lightblue4", + "penwidth": 3, + } + if edge_attrs is None + else edge_attrs + ) + self._default_cluster_attrs: dict = ( + { + **_default_attrs, + "shape": "rectangle", + "style": "solid", + } + if cluster_attrs is None + else cluster_attrs + ) + + def add_node( + self, + node_id: str, + node_label: str, + parent_graph_id: str | None = None, + **node_attrs: Any, + ) -> None: + """Add a single node to the graph. + + Args: + node_id (str): Unique node ID to identify this node. + node_label (str): The text to display on the node when rendered. + parent_graph_id (str | None): Optional ID of the cluster this node belongs to. + **node_attrs (Any): Any additional styling keyword arguments. + + """ + # Use ChainMap so you don't need to construct a new dictionary + node_attrs = ChainMap(node_attrs, self._default_node_attrs) + node = pydot.Node(node_id, label=node_label, **node_attrs) + parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id + + self._subgraphs[parent_graph_id].add_node(node) + + def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + """Add a single directed edge between nodes in the graph. + + Args: + from_node_id (str): The unique ID of the source node. + to_node_id (str): The unique ID of the destination node. + **edge_attrs (Any): Any additional styling keyword arguments. + + """ + # Use ChainMap so you don't need to construct a new dictionary + edge_attrs = ChainMap(edge_attrs, self._default_edge_attrs) + edge = pydot.Edge(from_node_id, to_node_id, **edge_attrs) + self.graph.add_edge(edge) + + def add_cluster( + self, + cluster_id: str, + cluster_label: str, + parent_graph_id: str | None = None, + **cluster_attrs: Any, + ) -> None: + """Add a single cluster to the graph. + + A cluster is a specific type of subgraph where the nodes and edges contained + within it are visually and logically grouped. + + Args: + cluster_id (str): Unique cluster ID to identify this cluster. + cluster_label (str): The text to display on the cluster when rendered. + parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. + **cluster_attrs (Any): Any additional styling keyword arguments. + + """ + # Use ChainMap so you don't need to construct a new dictionary + cluster_attrs = ChainMap(cluster_attrs, self._default_cluster_attrs) + cluster = pydot.Cluster(graph_name=cluster_id, **cluster_attrs) + + # Puts the label in a node within the cluster. + # Ensures that any edges connecting nodes through the cluster + # boundary don't block the label. + # ┌───────────┐ + # │ ┌───────┐ │ + # │ │ label │ │ + # │ └───────┘ │ + # │ │ + # └───────────┘ + if cluster_label: + node_id = f"{cluster_id}_info_node" + rank_subgraph = pydot.Subgraph() + node = pydot.Node( + node_id, + label=cluster_label, + shape="rectangle", + style="dashed", + fontname="Helvetica", + penwidth=2, + ) + rank_subgraph.add_node(node) + cluster.add_subgraph(rank_subgraph) + cluster.add_node(node) + + self._subgraphs[cluster_id] = cluster + + parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id + self._subgraphs[parent_graph_id].add_subgraph(cluster) + + def to_file(self, output_filename: str) -> None: + """Save the graph to a file. + + This method will infer the file's format (e.g., 'png', 'svg') from this filename's extension. + If no extension is provided, the 'png' format will be the default. + + Args: + output_filename (str): Desired filename for the graph. File extension can be included + and if no file extension is provided, it will default to a `.png` file. + + """ + output_filename_path: pathlib.Path = pathlib.Path(output_filename) + if not output_filename_path.suffix: + output_filename_path = output_filename_path.with_suffix(".png") + + format = output_filename_path.suffix[1:].lower() + + self.graph.write(str(output_filename_path), format=format) + + def to_string(self) -> str: + """Return the graph as a string. + + This is typically used to get the graph's representation in a standard string format like DOT. + + Returns: + str: A string representation of the graph. + """ + return self.graph.to_string() diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py new file mode 100644 index 0000000000..3f70ae8159 --- /dev/null +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -0,0 +1,94 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the DAGBuilder abstract base class.""" + +from typing import Any + +import pytest + +pytestmark = pytest.mark.usefixtures("requires_xdsl") + +# pylint: disable=wrong-import-position +# This import needs to be after pytest in order to prevent ImportErrors +from catalyst.python_interface.visualization.dag_builder import DAGBuilder + + +def test_concrete_implementation_works(): + """Unit test for concrete implementation of abc.""" + + # pylint: disable=unused-argument + class ConcreteDAGBuilder(DAGBuilder): + """Concrete subclass of an ABC for testing purposes.""" + + def add_node( + self, + node_id: str, + node_label: str, + parent_graph_id: str | None = None, + **node_attrs: Any, + ) -> None: + return + + def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + return + + def add_cluster( + self, + cluster_id: str, + cluster_label: str, + parent_graph_id: str | None = None, + **cluster_attrs: Any, + ) -> None: + return + + def to_file(self, output_filename: str) -> None: + return + + def to_string(self) -> str: + return "test" + + dag_builder = ConcreteDAGBuilder() + # pylint: disable = assignment-from-none + node = dag_builder.add_node("0", "node0") + edge = dag_builder.add_edge("0", "1") + cluster = dag_builder.add_cluster("0", "cluster0") + render = dag_builder.to_file("test.png") + string = dag_builder.to_string() + + assert node is None + assert edge is None + assert cluster is None + assert render is None + assert string == "test" + + +def test_abc_cannot_be_instantiated(): + """Tests that the DAGBuilder ABC cannot be instantiated.""" + + with pytest.raises(TypeError, match="Can't instantiate abstract class"): + # pylint: disable=abstract-class-instantiated + DAGBuilder() + + +def test_incomplete_subclass(): + """Tests that an incomplete subclass will fail""" + + # pylint: disable=too-few-public-methods + class IncompleteDAGBuilder(DAGBuilder): + def add_node(self, *args, **kwargs): + pass + + with pytest.raises(TypeError, match="Can't instantiate abstract class"): + # pylint: disable=abstract-class-instantiated + IncompleteDAGBuilder() diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py new file mode 100644 index 0000000000..5c57975a12 --- /dev/null +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -0,0 +1,300 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the PyDotDAGBuilder subclass.""" + +from unittest.mock import MagicMock + +import pytest + +pydot = pytest.importorskip("pydot") +pytestmark = pytest.mark.usefixtures("requires_xdsl") +# pylint: disable=wrong-import-position +from catalyst.python_interface.visualization.pydot_dag_builder import PyDotDAGBuilder + + +@pytest.mark.unit +def test_initialization_defaults(): + """Tests the default graph attributes are as expected.""" + + dag_builder = PyDotDAGBuilder() + + assert isinstance(dag_builder.graph, pydot.Dot) + # Ensure it's a directed graph + assert dag_builder.graph.get_graph_type() == "digraph" + # Ensure that it flows top to bottom + assert dag_builder.graph.get_rankdir() == "TB" + # Ensure edges can be connected directly to clusters / subgraphs + assert dag_builder.graph.get_compound() == "true" + # Ensure duplicated edges cannot be added + assert dag_builder.graph.obj_dict["strict"] is True + + +class TestAddMethods: + """Test that elements can be added to the graph.""" + + @pytest.mark.unit + def test_add_node(self): + """Unit test the `add_node` method.""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + node_list = dag_builder.graph.get_node_list() + assert len(node_list) == 1 + assert node_list[0].get_label() == "node0" + + @pytest.mark.unit + def test_add_edge(self): + """Unit test the `add_edge` method.""" + + dag_builder = PyDotDAGBuilder() + dag_builder.add_node("0", "node0") + dag_builder.add_node("1", "node1") + dag_builder.add_edge("0", "1") + + assert len(dag_builder.graph.get_edges()) == 1 + edge = dag_builder.graph.get_edges()[0] + assert edge.get_source() == "0" + assert edge.get_destination() == "1" + + @pytest.mark.unit + def test_add_cluster(self): + """Unit test the 'add_cluster' method.""" + + dag_builder = PyDotDAGBuilder() + dag_builder.add_cluster("0", "my_cluster") + + assert len(dag_builder.graph.get_subgraphs()) == 1 + assert dag_builder.graph.get_subgraphs()[0].get_name() == "cluster_0" + + @pytest.mark.unit + def test_add_node_to_parent_graph(self): + """Tests that you can add a node to a parent graph.""" + dag_builder = PyDotDAGBuilder() + + # Create node + dag_builder.add_node("0", "node0") + + # Create cluster + dag_builder.add_cluster("c0", "cluster0") + + # Create node inside cluster + dag_builder.add_node("1", "node1", parent_graph_id="c0") + + # Verify graph structure + root_graph = dag_builder.graph + + # Make sure the base graph has node0 + assert root_graph.get_node("0"), "Node 0 not found in root graph" + + # Get the cluster and verify it has node1 and not node0 + cluster_list = root_graph.get_subgraph("cluster_c0") + assert cluster_list, "Subgraph 'cluster_c0' not found" + cluster_graph = cluster_list[0] # Get the actual subgraph object + + assert cluster_graph.get_node("1"), "Node 1 not found in cluster 'c0'" + assert not cluster_graph.get_node("0"), ( + "Node 0 was incorrectly added to cluster" + ) + + assert not root_graph.get_node("1"), "Node 1 was incorrectly added to root" + + @pytest.mark.unit + def test_add_cluster_to_parent_graph(self): + """Test that you can add a cluster to a parent graph.""" + dag_builder = PyDotDAGBuilder() + + # Level 0 (Root): Adds cluster on top of base graph + dag_builder.add_node("n_root", "node_root") + dag_builder.add_cluster("c0", "cluster_outer") + + # Level 1 (Inside c0): Add node on outer cluster and create new cluster on top + dag_builder.add_node("n_outer", "node_outer", parent_graph_id="c0") + dag_builder.add_cluster("c1", "cluster_inner", parent_graph_id="c0") + + # Level 2 (Inside c1): Add node on second cluster + dag_builder.add_node("n_inner", "node_inner", parent_graph_id="c1") + + root_graph = dag_builder.graph + + outer_cluster_list = root_graph.get_subgraph("cluster_c0") + assert outer_cluster_list, "Outer cluster 'c0' not found in root" + c0 = outer_cluster_list[0] + + inner_cluster_list = c0.get_subgraph("cluster_c1") + assert inner_cluster_list, "Inner cluster 'c1' not found in 'c0'" + c1 = inner_cluster_list[0] + + # Check Level 0 (Root) + assert root_graph.get_node("n_root"), "n_root not found in root" + assert root_graph.get_subgraph("cluster_c0"), "c0 not found in root" + assert not root_graph.get_node("n_outer"), "n_outer incorrectly found in root" + assert not root_graph.get_node("n_inner"), "n_inner incorrectly found in root" + assert not root_graph.get_subgraph("cluster_c1"), "c1 incorrectly found in root" + + # Check Level 1 (c0) + assert c0.get_node("n_outer"), "n_outer not found in c0" + assert c0.get_subgraph("cluster_c1"), "c1 not found in c0" + assert not c0.get_node("n_root"), "n_root incorrectly found in c0" + assert not c0.get_node("n_inner"), "n_inner incorrectly found in c0" + + # Check Level 2 (c1) + assert c1.get_node("n_inner"), "n_inner not found in c1" + assert not c1.get_node("n_root"), "n_root incorrectly found in c1" + assert not c1.get_node("n_outer"), "n_outer incorrectly found in c1" + + +class TestAttributes: + """Tests that the attributes for elements in the graph are overridden correctly.""" + + @pytest.mark.unit + def test_default_graph_attrs(self): + """Test that default graph attributes can be set.""" + + dag_builder = PyDotDAGBuilder(attrs={"fontname": "Times"}) + + dag_builder.add_node("0", "node0") + node0 = dag_builder.graph.get_node("0")[0] + assert node0.get("fontname") == "Times" + + dag_builder.add_cluster("1", "cluster0") + cluster = dag_builder.graph.get_subgraphs()[0] + assert cluster.get("fontname") == "Times" + + @pytest.mark.unit + def test_add_node_with_attrs(self): + """Tests that default attributes are applied and can be overridden.""" + dag_builder = PyDotDAGBuilder( + node_attrs={"fillcolor": "lightblue", "penwidth": 3} + ) + + # Defaults + dag_builder.add_node("0", "node0") + node0 = dag_builder.graph.get_node("0")[0] + assert node0.get("fillcolor") == "lightblue" + assert node0.get("penwidth") == 3 + + # Make sure we can override + dag_builder.add_node("1", "node1", fillcolor="red", penwidth=4) + node1 = dag_builder.graph.get_node("1")[0] + assert node1.get("fillcolor") == "red" + assert node1.get("penwidth") == 4 + + @pytest.mark.unit + def test_add_edge_with_attrs(self): + """Tests that default attributes are applied and can be overridden.""" + dag_builder = PyDotDAGBuilder(edge_attrs={"color": "lightblue4", "penwidth": 3}) + + dag_builder.add_node("0", "node0") + dag_builder.add_node("1", "node1") + dag_builder.add_edge("0", "1") + edge = dag_builder.graph.get_edges()[0] + # Defaults defined earlier + assert edge.get("color") == "lightblue4" + assert edge.get("penwidth") == 3 + + # Make sure we can override + dag_builder.add_edge("0", "1", color="red", penwidth=4) + edge = dag_builder.graph.get_edges()[1] + assert edge.get("color") == "red" + assert edge.get("penwidth") == 4 + + @pytest.mark.unit + def test_add_cluster_with_attrs(self): + """Tests that default cluster attributes are applied and can be overridden.""" + dag_builder = PyDotDAGBuilder( + cluster_attrs={ + "style": "solid", + "fillcolor": None, + "penwidth": 2, + "fontname": "Helvetica", + } + ) + + dag_builder.add_cluster("0", "cluster0") + cluster1 = dag_builder.graph.get_subgraph("cluster_0")[0] + + # Defaults + assert cluster1.get("style") == "solid" + assert cluster1.get("fillcolor") is None + assert cluster1.get("penwidth") == 2 + assert cluster1.get("fontname") == "Helvetica" + + dag_builder.add_cluster( + "1", "cluster1", style="filled", penwidth=10, fillcolor="red" + ) + cluster2 = dag_builder.graph.get_subgraph("cluster_1")[0] + + # Make sure we can override + assert cluster2.get("style") == "filled" + assert cluster2.get("penwidth") == 10 + assert cluster2.get("fillcolor") == "red" + + # Check that other defaults are still present + assert cluster2.get("fontname") == "Helvetica" + + +class TestOutput: + """Test that the graph can be outputted correctly.""" + + @pytest.mark.unit + @pytest.mark.parametrize( + "filename, format", + [("my_graph", None), ("my_graph", "png"), ("prototype.trial1", "png")], + ) + def test_to_file(self, monkeypatch, filename, format): + """Tests that the `to_file` method works correctly.""" + dag_builder = PyDotDAGBuilder() + + # mock out the graph writing functionality + mock_write = MagicMock() + monkeypatch.setattr(dag_builder.graph, "write", mock_write) + dag_builder.to_file(filename + "." + (format or "png")) + + # make sure the function handles extensions correctly + mock_write.assert_called_once_with( + filename + "." + (format or "png"), format=format or "png" + ) + + @pytest.mark.unit + @pytest.mark.parametrize("format", ["pdf", "svg", "jpeg"]) + def test_other_supported_formats(self, monkeypatch, format): + """Tests that the `to_file` method works with other formats.""" + dag_builder = PyDotDAGBuilder() + + # mock out the graph writing functionality + mock_write = MagicMock() + monkeypatch.setattr(dag_builder.graph, "write", mock_write) + dag_builder.to_file(f"my_graph.{format}") + + # make sure the function handles extensions correctly + mock_write.assert_called_once_with(f"my_graph.{format}", format=format) + + @pytest.mark.unit + def test_to_string(self): + """Tests that the `to_string` method works correclty.""" + + dag_builder = PyDotDAGBuilder() + dag_builder.add_node("n0", "node0") + dag_builder.add_node("n1", "node1") + dag_builder.add_edge("n0", "n1") + + string = dag_builder.to_string() + assert isinstance(string, str) + + # make sure important things show up in the string + assert "digraph" in string + assert "n0" in string + assert "n1" in string + assert "n0 -> n1" in string From e9ceaa7764b6f26befcea6e2340647e76ac077df Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 13:33:20 -0500 Subject: [PATCH 002/245] add pydot to requirements.txt --- requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/requirements.txt b/requirements.txt index 08f57ee9c0..6b904318c6 100644 --- a/requirements.txt +++ b/requirements.txt @@ -36,3 +36,4 @@ pennylane-lightning-kokkos amazon-braket-pennylane-plugin>1.27.1 xdsl xdsl-jax +pydot From ae50ca166c47fa9e3343c665da8144593c375734 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 15:40:07 -0500 Subject: [PATCH 003/245] add files --- .../python_interface/visualization/mlir_dag_analysis_pass.py | 0 .../python_interface/visualization/test_mlir_dag_analysis_pass.py | 0 2 files changed, 0 insertions(+), 0 deletions(-) create mode 100644 frontend/catalyst/python_interface/visualization/mlir_dag_analysis_pass.py create mode 100644 frontend/test/pytest/python_interface/visualization/test_mlir_dag_analysis_pass.py diff --git a/frontend/catalyst/python_interface/visualization/mlir_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/mlir_dag_analysis_pass.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/frontend/test/pytest/python_interface/visualization/test_mlir_dag_analysis_pass.py b/frontend/test/pytest/python_interface/visualization/test_mlir_dag_analysis_pass.py new file mode 100644 index 0000000000..e69de29bb2 From a0e37b054ed3da037f049f41fdb8f9e436251aed Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 15:41:57 -0500 Subject: [PATCH 004/245] rename --- .../{mlir_dag_analysis_pass.py => circuit_dag_analysis_pass.py} | 0 ...lir_dag_analysis_pass.py => test_circuit_dag_analysis_pass.py} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename frontend/catalyst/python_interface/visualization/{mlir_dag_analysis_pass.py => circuit_dag_analysis_pass.py} (100%) rename frontend/test/pytest/python_interface/visualization/{test_mlir_dag_analysis_pass.py => test_circuit_dag_analysis_pass.py} (100%) diff --git a/frontend/catalyst/python_interface/visualization/mlir_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py similarity index 100% rename from frontend/catalyst/python_interface/visualization/mlir_dag_analysis_pass.py rename to frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py diff --git a/frontend/test/pytest/python_interface/visualization/test_mlir_dag_analysis_pass.py b/frontend/test/pytest/python_interface/visualization/test_circuit_dag_analysis_pass.py similarity index 100% rename from frontend/test/pytest/python_interface/visualization/test_mlir_dag_analysis_pass.py rename to frontend/test/pytest/python_interface/visualization/test_circuit_dag_analysis_pass.py From 76b55653c656e1f4b6cef10714adc680079c51ad Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 16:08:03 -0500 Subject: [PATCH 005/245] add base class --- .../circuit_dag_analysis_pass.py | 61 +++++++++++++++++++ 1 file changed, 61 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index e69de29bb2..8baf52dd23 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -0,0 +1,61 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contains the CircuitDAGAnalysisPass for generating a DAG from an xDSL module.""" + +from functools import singledispatchmethod +from typing import TYPE_CHECKING, Any + +import xdsl +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Region + +from catalyst.python_interface.dialects import quantum + +if TYPE_CHECKING: + from catalyst.python_interface.visualization.dag_builder import DAGBuilder + + +class CircuitDAGAnalysisPass: + def __init__(self, dag_builder: DAGBuilder) -> None: + """Initialize the analysis pass.""" + self.dag_builder: DAGBuilder = dag_builder + + @singledispatchmethod + def visit_op(self, op: Any) -> None: + """Default handler for unknown operation types. + + This method is dispatched based on the type of 'op'. + + Args: + op (Any): An xDSL operation. + """ + + @visit_op.register + def visit_region(self, region: Region) -> None: + """Visit an xDSL Region operation.""" + for block in region.blocks: + self.visit_block(block) + + @visit_op.register + def visit_block(self, block: Block) -> None: + """Visit an xDSL Block operation.""" + for op in block.ops: + self.visit_op(op) + + def run(self, module: builtin.ModuleOp) -> None: + """Apply the analysis pass on the module.""" + + for op in module.ops: + self.visit_op(op) From 659f480894841b3069bb21947a6c8acb517ab498 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 16:26:36 -0500 Subject: [PATCH 006/245] quick changes --- .../visualization/circuit_dag_analysis_pass.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index 8baf52dd23..d439a422ca 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -22,9 +22,7 @@ from xdsl.ir import Block, Region from catalyst.python_interface.dialects import quantum - -if TYPE_CHECKING: - from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from catalyst.python_interface.visualization.dag_builder import DAGBuilder class CircuitDAGAnalysisPass: From 8699557dc2d25b37b98477947db5f4fef0cf911a Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 16:42:33 -0500 Subject: [PATCH 007/245] add control flow support --- .../visualization/circuit_dag_analysis_pass.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index d439a422ca..f5b63537cb 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -40,6 +40,23 @@ def visit_op(self, op: Any) -> None: op (Any): An xDSL operation. """ + # ╔═══════════════════════════════════════════════════════════════╗ + # ║ CONTROL FLOW HANDLERS: Specialized dispatch for xDSL control  ║ + # ║                         flow operations (scf dialect).        ║ + # ╚═══════════════════════════════════════════════════════════════╝ + + @visit_op.register + def _visit_for_op(self, op: scf.ForOp) -> None: + """Handle an xDSL ForOp operation.""" + + @visit_op.register + def _visit_while_op(self, op: scf.WhileOp) -> None: + """Handle an xDSL WhileOp operation.""" + + @visit_op.register + def _visit_if_op(self, op: scf.IfOp) -> None: + """Handle an xDSL WhileOp operation.""" + @visit_op.register def visit_region(self, region: Region) -> None: """Visit an xDSL Region operation.""" From db6edc9f0397f220bdbbd15016cf707749cfe692 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 19 Nov 2025 17:04:15 -0500 Subject: [PATCH 008/245] clean-up --- .../circuit_dag_analysis_pass.py | 111 +++++++++++++----- 1 file changed, 83 insertions(+), 28 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index f5b63537cb..bbe8fc8294 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -26,51 +26,106 @@ class CircuitDAGAnalysisPass: - def __init__(self, dag_builder: DAGBuilder) -> None: - """Initialize the analysis pass.""" - self.dag_builder: DAGBuilder = dag_builder - - @singledispatchmethod - def visit_op(self, op: Any) -> None: - """Default handler for unknown operation types. + """A Pass that analyzes an xDSL module and constructs a Directed Acyclic Graph (DAG) + using an injected DAGBuilder instance. This is a non-mutating Analysis Pass.""" - This method is dispatched based on the type of 'op'. + def __init__(self, dag_builder: DAGBuilder) -> None: + """Initialize the analysis pass by injecting the DAG builder dependency. Args: - op (Any): An xDSL operation. + dag_builder (DAGBuilder): The concrete builder instance used for graph construction. """ + self.dag_builder: DAGBuilder = dag_builder - # ╔═══════════════════════════════════════════════════════════════╗ - # ║ CONTROL FLOW HANDLERS: Specialized dispatch for xDSL control  ║ - # ║                         flow operations (scf dialect).        ║ - # ╚═══════════════════════════════════════════════════════════════╝ + # ================================= + # 1. CORE DISPATCH AND ENTRY POINT + # ================================= - @visit_op.register - def _visit_for_op(self, op: scf.ForOp) -> None: - """Handle an xDSL ForOp operation.""" + @singledispatchmethod + def visit_op(self, op: Any) -> None: + """Central dispatch method (Visitor Pattern). Routes the operation 'op' + to the specialized handler registered for its type.""" + pass - @visit_op.register - def _visit_while_op(self, op: scf.WhileOp) -> None: - """Handle an xDSL WhileOp operation.""" + def run(self, module: builtin.ModuleOp) -> None: + """Applies the analysis pass on the module.""" + for op in module.ops: + self.visit_op(op) - @visit_op.register - def _visit_if_op(self, op: scf.IfOp) -> None: - """Handle an xDSL WhileOp operation.""" + # ======================= + # 2. HIERARCHY TRAVERSAL + # ======================= + # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). @visit_op.register def visit_region(self, region: Region) -> None: - """Visit an xDSL Region operation.""" + """Visit an xDSL Region operation, delegating traversal to its Blocks.""" for block in region.blocks: self.visit_block(block) @visit_op.register def visit_block(self, block: Block) -> None: - """Visit an xDSL Block operation.""" + """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: self.visit_op(op) - def run(self, module: builtin.ModuleOp) -> None: - """Apply the analysis pass on the module.""" + # ====================================== + # 3. QUANTUM GATE & STATE PREP HANDLERS + # ====================================== + # Handlers for operations that apply unitary transformations or set-up the quantum state. - for op in module.ops: - self.visit_op(op) + @visit_op.register + def _visit_unitary_and_state_prep( + self, + op: ( + quantum.CustomOp + | quantum.GlobalPhaseOp + | quantum.QubitUnitaryOp + | quantum.MultiRZOp + | quantum.SetStateOp + | quantum.SetBasisStateOp + ), + ) -> None: + """Generic handler for unitary gates and quantum state preparation operations.""" + pass + + # ============================================= + # 4. QUANTUM MEASUREMENT & OBSERVABLE HANDLERS + # ============================================= + + @visit_op.register + def _visit_terminal_state_op(self, op: quantum.StateOp) -> None: + """Handler for the terminal StateOp, which retrieves the final state vector.""" + pass + + @visit_op.register + def _visit_statistical_measurement_ops( + self, + op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, + ) -> None: + """Handler for statistical measurement operations (e.g., Expval, Sample).""" + pass + + @visit_op.register + def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: + """Handler for the single-qubit projective MeasureOp.""" + pass + + # ========================= + # 5. CONTROL FLOW HANDLERS + # ========================= + + @visit_op.register + def _visit_for_op(self, op: scf.ForOp) -> None: + """Handle an xDSL ForOp operation (Loop cluster creation).""" + pass + + @visit_op.register + def _visit_while_op(self, op: scf.WhileOp) -> None: + """Handle an xDSL WhileOp operation (Loop cluster creation).""" + pass + + @visit_op.register + def _visit_if_op(self, op: scf.IfOp) -> None: + """Handle an xDSL IfOp operation (Conditional cluster creation).""" + pass From 2be63e496425f9ab3d84cc4e5ed46a04492ebf6a Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 19 Nov 2025 19:20:20 -0500 Subject: [PATCH 009/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/circuit_dag_analysis_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index bbe8fc8294..bb93239dae 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -94,7 +94,7 @@ def _visit_unitary_and_state_prep( # ============================================= @visit_op.register - def _visit_terminal_state_op(self, op: quantum.StateOp) -> None: + def _visit_state_op(self, op: quantum.StateOp) -> None: """Handler for the terminal StateOp, which retrieves the final state vector.""" pass From 12fa2afd08ee2f0e735c263bb864943d668259ec Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 19 Nov 2025 19:20:33 -0500 Subject: [PATCH 010/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/circuit_dag_analysis_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index bb93239dae..6e80816d76 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -95,7 +95,7 @@ def _visit_unitary_and_state_prep( @visit_op.register def _visit_state_op(self, op: quantum.StateOp) -> None: - """Handler for the terminal StateOp, which retrieves the final state vector.""" + """Handler for the terminal StateOp.""" pass @visit_op.register From 2c933309a70ebd7be53060d42ff85660c156908a Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 19 Nov 2025 19:21:04 -0500 Subject: [PATCH 011/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/circuit_dag_analysis_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index 6e80816d76..891abf909b 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -103,7 +103,7 @@ def _visit_statistical_measurement_ops( self, op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, ) -> None: - """Handler for statistical measurement operations (e.g., Expval, Sample).""" + """Handler for statistical measurement operations.""" pass @visit_op.register From e68c3f434b7e0bf2be6f2607fad93a0c4718ce01 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 19 Nov 2025 19:21:34 -0500 Subject: [PATCH 012/245] Apply suggestion from @andrijapau --- .../visualization/circuit_dag_analysis_pass.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index 891abf909b..d4369ff617 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -117,15 +117,15 @@ def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: @visit_op.register def _visit_for_op(self, op: scf.ForOp) -> None: - """Handle an xDSL ForOp operation (Loop cluster creation).""" + """Handle an xDSL ForOp operation.""" pass @visit_op.register def _visit_while_op(self, op: scf.WhileOp) -> None: - """Handle an xDSL WhileOp operation (Loop cluster creation).""" + """Handle an xDSL WhileOp operation.""" pass @visit_op.register def _visit_if_op(self, op: scf.IfOp) -> None: - """Handle an xDSL IfOp operation (Conditional cluster creation).""" + """Handle an xDSL IfOp operation.""" pass From 4aad864ec6dfc912ba2b2716be6304575fd7334b Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 19 Nov 2025 19:22:01 -0500 Subject: [PATCH 013/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/circuit_dag_analysis_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index d4369ff617..528ac54d95 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -59,7 +59,7 @@ def run(self, module: builtin.ModuleOp) -> None: @visit_op.register def visit_region(self, region: Region) -> None: - """Visit an xDSL Region operation, delegating traversal to its Blocks.""" + """Visit an xDSL Region operation.""" for block in region.blocks: self.visit_block(block) From 3be39dcbfa0fcf99fbb90761844673d5a5d4fec0 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 11:23:29 -0500 Subject: [PATCH 014/245] update cluster label logic --- .../python_interface/visualization/dag_builder.py | 4 ++-- .../python_interface/visualization/pydot_dag_builder.py | 8 ++++---- 2 files changed, 6 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 7a60a55f1e..5526ee5bf7 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -57,7 +57,7 @@ def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> Non def add_cluster( self, cluster_id: str, - cluster_label: str, + node_label: str | None = None, parent_graph_id: str | None = None, **cluster_attrs: Any, ) -> None: @@ -68,7 +68,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. - cluster_label (str): The text to display on the cluster when rendered. + node_label (str): The text to display on an information node within the cluster when rendered. parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. **cluster_attrs (Any): Any additional styling keyword arguments. diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index ad9eef4bcb..f52e0fa834 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -128,7 +128,7 @@ def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> Non def add_cluster( self, cluster_id: str, - cluster_label: str, + node_label: str | None = None, parent_graph_id: str | None = None, **cluster_attrs: Any, ) -> None: @@ -139,7 +139,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. - cluster_label (str): The text to display on the cluster when rendered. + node_label (str): The text to display on the information node on the cluster when rendered. parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. **cluster_attrs (Any): Any additional styling keyword arguments. @@ -157,12 +157,12 @@ def add_cluster( # │ └───────┘ │ # │ │ # └───────────┘ - if cluster_label: + if node_label: node_id = f"{cluster_id}_info_node" rank_subgraph = pydot.Subgraph() node = pydot.Node( node_id, - label=cluster_label, + label=node_label, shape="rectangle", style="dashed", fontname="Helvetica", From a383ba8fb26f655d8b86ccb4149651dda5627a0e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 11:25:20 -0500 Subject: [PATCH 015/245] fix dag builders test --- .../python_interface/visualization/test_dag_builder.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py index 3f70ae8159..2e935bae1b 100644 --- a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -40,13 +40,15 @@ def add_node( ) -> None: return - def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + def add_edge( + self, from_node_id: str, to_node_id: str, **edge_attrs: Any + ) -> None: return def add_cluster( self, cluster_id: str, - cluster_label: str, + node_label: str | None = None, parent_graph_id: str | None = None, **cluster_attrs: Any, ) -> None: @@ -62,7 +64,7 @@ def to_string(self) -> str: # pylint: disable = assignment-from-none node = dag_builder.add_node("0", "node0") edge = dag_builder.add_edge("0", "1") - cluster = dag_builder.add_cluster("0", "cluster0") + cluster = dag_builder.add_cluster("0") render = dag_builder.to_file("test.png") string = dag_builder.to_string() From 8ddaa5d36a7945e18233c18179474db6c2d63d7d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 11:26:07 -0500 Subject: [PATCH 016/245] fix pydot dag builders test --- .../visualization/test_pydot_dag_builder.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 5c57975a12..11e5e4ac47 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -73,7 +73,7 @@ def test_add_cluster(self): """Unit test the 'add_cluster' method.""" dag_builder = PyDotDAGBuilder() - dag_builder.add_cluster("0", "my_cluster") + dag_builder.add_cluster("0") assert len(dag_builder.graph.get_subgraphs()) == 1 assert dag_builder.graph.get_subgraphs()[0].get_name() == "cluster_0" @@ -87,7 +87,7 @@ def test_add_node_to_parent_graph(self): dag_builder.add_node("0", "node0") # Create cluster - dag_builder.add_cluster("c0", "cluster0") + dag_builder.add_cluster("c0") # Create node inside cluster dag_builder.add_node("1", "node1", parent_graph_id="c0") @@ -117,11 +117,11 @@ def test_add_cluster_to_parent_graph(self): # Level 0 (Root): Adds cluster on top of base graph dag_builder.add_node("n_root", "node_root") - dag_builder.add_cluster("c0", "cluster_outer") + dag_builder.add_cluster("c0") # Level 1 (Inside c0): Add node on outer cluster and create new cluster on top dag_builder.add_node("n_outer", "node_outer", parent_graph_id="c0") - dag_builder.add_cluster("c1", "cluster_inner", parent_graph_id="c0") + dag_builder.add_cluster("c1", parent_graph_id="c0") # Level 2 (Inside c1): Add node on second cluster dag_builder.add_node("n_inner", "node_inner", parent_graph_id="c1") @@ -168,7 +168,7 @@ def test_default_graph_attrs(self): node0 = dag_builder.graph.get_node("0")[0] assert node0.get("fontname") == "Times" - dag_builder.add_cluster("1", "cluster0") + dag_builder.add_cluster("1") cluster = dag_builder.graph.get_subgraphs()[0] assert cluster.get("fontname") == "Times" @@ -222,7 +222,7 @@ def test_add_cluster_with_attrs(self): } ) - dag_builder.add_cluster("0", "cluster0") + dag_builder.add_cluster("0") cluster1 = dag_builder.graph.get_subgraph("cluster_0")[0] # Defaults @@ -232,7 +232,7 @@ def test_add_cluster_with_attrs(self): assert cluster1.get("fontname") == "Helvetica" dag_builder.add_cluster( - "1", "cluster1", style="filled", penwidth=10, fillcolor="red" + "1", style="filled", penwidth=10, fillcolor="red" ) cluster2 = dag_builder.graph.get_subgraph("cluster_1")[0] From 1166f40c8d4bb87e5e61c306b19fd5a6eb94aed9 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 11:27:41 -0500 Subject: [PATCH 017/245] update doc --- frontend/catalyst/python_interface/visualization/dag_builder.py | 2 +- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 5526ee5bf7..aeb8b217b6 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -68,7 +68,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. - node_label (str): The text to display on an information node within the cluster when rendered. + node_label (str | None): The text to display on an information node within the cluster when rendered. parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. **cluster_attrs (Any): Any additional styling keyword arguments. diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index f52e0fa834..cff97b0148 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -139,7 +139,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. - node_label (str): The text to display on the information node on the cluster when rendered. + node_label (str | None): The text to display on the information node on the cluster when rendered. parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. **cluster_attrs (Any): Any additional styling keyword arguments. From df9bf45e6648a0127a2d4cb356a2c4f10e36a1a9 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 12:33:00 -0500 Subject: [PATCH 018/245] basic cl --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 831453990c..abb91a6bdd 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -3,7 +3,7 @@

New features since last release

* Compiled programs can be visualized. - [(#)]() + [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From 8ca5b1fe925e03dbb405d1b0b657c3e87f773e84 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 12:33:35 -0500 Subject: [PATCH 019/245] basic cl --- doc/releases/changelog-dev.md | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index abb91a6bdd..3e29074a66 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -4,6 +4,7 @@ * Compiled programs can be visualized. [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213) + [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From b579a3fe57d9f20022a2a4387432b8a8a6802111 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 12:58:22 -0500 Subject: [PATCH 020/245] Trigger CI From a94d8432054cdab46bfb94702675f49f88b755b9 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 12:58:28 -0500 Subject: [PATCH 021/245] Trigger CI From 331e3491404c75d67a8cec30a287f4bb7f2b8618 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 20 Nov 2025 13:09:45 -0500 Subject: [PATCH 022/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/circuit_dag_analysis_pass.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index 528ac54d95..df1a316e0a 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -90,7 +90,7 @@ def _visit_unitary_and_state_prep( pass # ============================================= - # 4. QUANTUM MEASUREMENT & OBSERVABLE HANDLERS + # 4. QUANTUM MEASUREMENT HANDLERS # ============================================= @visit_op.register From a87d11744f3d8c307573e8f0b65f54784ae57cc4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 13:12:41 -0500 Subject: [PATCH 023/245] just do customop --- .../visualization/circuit_dag_analysis_pass.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index df1a316e0a..7fa5d71065 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -77,14 +77,7 @@ def visit_block(self, block: Block) -> None: @visit_op.register def _visit_unitary_and_state_prep( self, - op: ( - quantum.CustomOp - | quantum.GlobalPhaseOp - | quantum.QubitUnitaryOp - | quantum.MultiRZOp - | quantum.SetStateOp - | quantum.SetBasisStateOp - ), + op: quantum.CustomOp, ) -> None: """Generic handler for unitary gates and quantum state preparation operations.""" pass From 9713aa3fa1c2150c88673c48dd5872f8afa4e906 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 13:15:22 -0500 Subject: [PATCH 024/245] fix wording --- .../visualization/circuit_dag_analysis_pass.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py index 7fa5d71065..bcc229798b 100644 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py @@ -45,7 +45,7 @@ def __init__(self, dag_builder: DAGBuilder) -> None: def visit_op(self, op: Any) -> None: """Central dispatch method (Visitor Pattern). Routes the operation 'op' to the specialized handler registered for its type.""" - pass + raise NotImplementedError(f"Dispatch not registered for operator of type {type(op)}") def run(self, module: builtin.ModuleOp) -> None: """Applies the analysis pass on the module.""" @@ -88,7 +88,7 @@ def _visit_unitary_and_state_prep( @visit_op.register def _visit_state_op(self, op: quantum.StateOp) -> None: - """Handler for the terminal StateOp.""" + """Handler for the terminal state measurement operation.""" pass @visit_op.register @@ -101,7 +101,7 @@ def _visit_statistical_measurement_ops( @visit_op.register def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: - """Handler for the single-qubit projective MeasureOp.""" + """Handler for the single-qubit projective measurement operation.""" pass # ========================= From 7bbc22410d50c89b83d6e81aaa4c92ac340ec924 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 14:13:24 -0500 Subject: [PATCH 025/245] rename --- .../circuit_dag_analysis_pass.py | 124 ------------------ .../test_circuit_dag_analysis_pass.py | 0 2 files changed, 124 deletions(-) delete mode 100644 frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py delete mode 100644 frontend/test/pytest/python_interface/visualization/test_circuit_dag_analysis_pass.py diff --git a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py b/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py deleted file mode 100644 index bcc229798b..0000000000 --- a/frontend/catalyst/python_interface/visualization/circuit_dag_analysis_pass.py +++ /dev/null @@ -1,124 +0,0 @@ -# Copyright 2025 Xanadu Quantum Technologies Inc. - -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at - -# http://www.apache.org/licenses/LICENSE-2.0 - -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -"""Contains the CircuitDAGAnalysisPass for generating a DAG from an xDSL module.""" - -from functools import singledispatchmethod -from typing import TYPE_CHECKING, Any - -import xdsl -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Region - -from catalyst.python_interface.dialects import quantum -from catalyst.python_interface.visualization.dag_builder import DAGBuilder - - -class CircuitDAGAnalysisPass: - """A Pass that analyzes an xDSL module and constructs a Directed Acyclic Graph (DAG) - using an injected DAGBuilder instance. This is a non-mutating Analysis Pass.""" - - def __init__(self, dag_builder: DAGBuilder) -> None: - """Initialize the analysis pass by injecting the DAG builder dependency. - - Args: - dag_builder (DAGBuilder): The concrete builder instance used for graph construction. - """ - self.dag_builder: DAGBuilder = dag_builder - - # ================================= - # 1. CORE DISPATCH AND ENTRY POINT - # ================================= - - @singledispatchmethod - def visit_op(self, op: Any) -> None: - """Central dispatch method (Visitor Pattern). Routes the operation 'op' - to the specialized handler registered for its type.""" - raise NotImplementedError(f"Dispatch not registered for operator of type {type(op)}") - - def run(self, module: builtin.ModuleOp) -> None: - """Applies the analysis pass on the module.""" - for op in module.ops: - self.visit_op(op) - - # ======================= - # 2. HIERARCHY TRAVERSAL - # ======================= - # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). - - @visit_op.register - def visit_region(self, region: Region) -> None: - """Visit an xDSL Region operation.""" - for block in region.blocks: - self.visit_block(block) - - @visit_op.register - def visit_block(self, block: Block) -> None: - """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" - for op in block.ops: - self.visit_op(op) - - # ====================================== - # 3. QUANTUM GATE & STATE PREP HANDLERS - # ====================================== - # Handlers for operations that apply unitary transformations or set-up the quantum state. - - @visit_op.register - def _visit_unitary_and_state_prep( - self, - op: quantum.CustomOp, - ) -> None: - """Generic handler for unitary gates and quantum state preparation operations.""" - pass - - # ============================================= - # 4. QUANTUM MEASUREMENT HANDLERS - # ============================================= - - @visit_op.register - def _visit_state_op(self, op: quantum.StateOp) -> None: - """Handler for the terminal state measurement operation.""" - pass - - @visit_op.register - def _visit_statistical_measurement_ops( - self, - op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, - ) -> None: - """Handler for statistical measurement operations.""" - pass - - @visit_op.register - def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: - """Handler for the single-qubit projective measurement operation.""" - pass - - # ========================= - # 5. CONTROL FLOW HANDLERS - # ========================= - - @visit_op.register - def _visit_for_op(self, op: scf.ForOp) -> None: - """Handle an xDSL ForOp operation.""" - pass - - @visit_op.register - def _visit_while_op(self, op: scf.WhileOp) -> None: - """Handle an xDSL WhileOp operation.""" - pass - - @visit_op.register - def _visit_if_op(self, op: scf.IfOp) -> None: - """Handle an xDSL IfOp operation.""" - pass diff --git a/frontend/test/pytest/python_interface/visualization/test_circuit_dag_analysis_pass.py b/frontend/test/pytest/python_interface/visualization/test_circuit_dag_analysis_pass.py deleted file mode 100644 index e69de29bb2..0000000000 From d6db96550c425f8f0b5e33e4f6484b8077fc7e82 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 14:14:07 -0500 Subject: [PATCH 026/245] add back --- .../visualization/construct_circuit_dag.py | 124 ++++++++++++++++++ .../test_construct_circuit_dag.py | 0 2 files changed, 124 insertions(+) create mode 100644 frontend/catalyst/python_interface/visualization/construct_circuit_dag.py create mode 100644 frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py new file mode 100644 index 0000000000..1237f94fed --- /dev/null +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -0,0 +1,124 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module.""" + +from functools import singledispatchmethod +from typing import TYPE_CHECKING, Any + +import xdsl +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Region + +from catalyst.python_interface.dialects import quantum +from catalyst.python_interface.visualization.dag_builder import DAGBuilder + + +class ConstructCircuitDAG: + """A tool that analyzes an xDSL module and constructs a Directed Acyclic Graph (DAG) + using an injected DAGBuilder instance. This tool does not mutate the xDSL module.""" + + def __init__(self, dag_builder: DAGBuilder) -> None: + """Initialize the analysis pass by injecting the DAG builder dependency. + + Args: + dag_builder (DAGBuilder): The concrete builder instance used for graph construction. + """ + self.dag_builder: DAGBuilder = dag_builder + + # ================================= + # 1. CORE DISPATCH AND ENTRY POINT + # ================================= + + @singledispatchmethod + def visit_op(self, op: Any) -> None: + """Central dispatch method (Visitor Pattern). Routes the operation 'op' + to the specialized handler registered for its type.""" + raise NotImplementedError(f"Dispatch not registered for operator of type {type(op)}") + + def construct(self, module: builtin.ModuleOp) -> None: + """Constructs the DAG from the module.""" + for op in module.ops: + self.visit_op(op) + + # ======================= + # 2. HIERARCHY TRAVERSAL + # ======================= + # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). + + @visit_op.register + def visit_region(self, region: Region) -> None: + """Visit an xDSL Region operation.""" + for block in region.blocks: + self.visit_block(block) + + @visit_op.register + def visit_block(self, block: Block) -> None: + """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" + for op in block.ops: + self.visit_op(op) + + # ====================================== + # 3. QUANTUM GATE & STATE PREP HANDLERS + # ====================================== + # Handlers for operations that apply unitary transformations or set-up the quantum state. + + @visit_op.register + def _visit_unitary_and_state_prep( + self, + op: quantum.CustomOp, + ) -> None: + """Generic handler for unitary gates and quantum state preparation operations.""" + pass + + # ============================================= + # 4. QUANTUM MEASUREMENT HANDLERS + # ============================================= + + @visit_op.register + def _visit_state_op(self, op: quantum.StateOp) -> None: + """Handler for the terminal state measurement operation.""" + pass + + @visit_op.register + def _visit_statistical_measurement_ops( + self, + op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, + ) -> None: + """Handler for statistical measurement operations.""" + pass + + @visit_op.register + def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: + """Handler for the single-qubit projective measurement operation.""" + pass + + # ========================= + # 5. CONTROL FLOW HANDLERS + # ========================= + + @visit_op.register + def _visit_for_op(self, op: scf.ForOp) -> None: + """Handle an xDSL ForOp operation.""" + pass + + @visit_op.register + def _visit_while_op(self, op: scf.WhileOp) -> None: + """Handle an xDSL WhileOp operation.""" + pass + + @visit_op.register + def _visit_if_op(self, op: scf.IfOp) -> None: + """Handle an xDSL IfOp operation.""" + pass diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py new file mode 100644 index 0000000000..e69de29bb2 From 11b08f026682398898550161c2a31f6a3bac6015 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 20 Nov 2025 14:53:22 -0500 Subject: [PATCH 027/245] Update frontend/catalyst/python_interface/visualization/pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index cff97b0148..4c7c40a579 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -49,7 +49,7 @@ def __init__( # - graph_type="digraph": Create a directed graph (edges have arrows). # - rankdir="TB": Set layout direction from Top to Bottom. # - compound="true": Allow edges to connect directly to clusters/subgraphs. - # - strict=True": Prevent duplicate edges (e.g., A -> B added twice). + # - strict=True: Prevent duplicate edges (e.g., A -> B added twice). self.graph: pydot.Dot = pydot.Dot( graph_type="digraph", rankdir="TB", compound="true", strict=True ) From a76b5bdfc4efb9f7fef884255d60c757f0803e4e Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 20 Nov 2025 14:53:27 -0500 Subject: [PATCH 028/245] Update frontend/catalyst/python_interface/visualization/pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 4c7c40a579..384537c93e 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -139,7 +139,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. - node_label (str | None): The text to display on the information node on the cluster when rendered. + node_label (str | None): The text to display on the information node within the cluster when rendered. parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. **cluster_attrs (Any): Any additional styling keyword arguments. From 90e4ebd4e7281deab5ad1615e94474384c777ec5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 15:16:24 -0500 Subject: [PATCH 029/245] fix --- .../visualization/construct_circuit_dag.py | 48 +++++++++++-------- 1 file changed, 27 insertions(+), 21 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 1237f94fed..27b7a2c631 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -19,7 +19,7 @@ import xdsl from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Region +from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder @@ -42,40 +42,46 @@ def __init__(self, dag_builder: DAGBuilder) -> None: # ================================= @singledispatchmethod - def visit_op(self, op: Any) -> None: + def visit(self, op: Any) -> None: """Central dispatch method (Visitor Pattern). Routes the operation 'op' to the specialized handler registered for its type.""" - raise NotImplementedError(f"Dispatch not registered for operator of type {type(op)}") + pass def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module.""" for op in module.ops: - self.visit_op(op) + self.visit(op) # ======================= # 2. HIERARCHY TRAVERSAL # ======================= # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). - @visit_op.register + @visit.register + def visit_operation(self, operation: Operation) -> None: + """Visit an xDSL Operation.""" + for region in operation.regions: + self.visit_region(region) + + @visit.register def visit_region(self, region: Region) -> None: """Visit an xDSL Region operation.""" for block in region.blocks: self.visit_block(block) - @visit_op.register + @visit.register def visit_block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: - self.visit_op(op) + self.visit(op) # ====================================== # 3. QUANTUM GATE & STATE PREP HANDLERS # ====================================== # Handlers for operations that apply unitary transformations or set-up the quantum state. - @visit_op.register - def _visit_unitary_and_state_prep( + @visit.register + def _unitary_and_state_prep( self, op: quantum.CustomOp, ) -> None: @@ -86,21 +92,21 @@ def _visit_unitary_and_state_prep( # 4. QUANTUM MEASUREMENT HANDLERS # ============================================= - @visit_op.register - def _visit_state_op(self, op: quantum.StateOp) -> None: + @visit.register + def _state_op(self, op: quantum.StateOp) -> None: """Handler for the terminal state measurement operation.""" pass - @visit_op.register - def _visit_statistical_measurement_ops( + @visit.register + def _statistical_measurement_ops( self, op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, ) -> None: """Handler for statistical measurement operations.""" pass - @visit_op.register - def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: + @visit.register + def _projective_measure_op(self, op: quantum.MeasureOp) -> None: """Handler for the single-qubit projective measurement operation.""" pass @@ -108,17 +114,17 @@ def _visit_projective_measure_op(self, op: quantum.MeasureOp) -> None: # 5. CONTROL FLOW HANDLERS # ========================= - @visit_op.register - def _visit_for_op(self, op: scf.ForOp) -> None: + @visit.register + def _for_op(self, op: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" pass - @visit_op.register - def _visit_while_op(self, op: scf.WhileOp) -> None: + @visit.register + def _while_op(self, op: scf.WhileOp) -> None: """Handle an xDSL WhileOp operation.""" pass - @visit_op.register - def _visit_if_op(self, op: scf.IfOp) -> None: + @visit.register + def _if_op(self, op: scf.IfOp) -> None: """Handle an xDSL IfOp operation.""" pass From 3d9e4bb25387631a0adc943958110ef4f2efc9ff Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 20 Nov 2025 15:16:54 -0500 Subject: [PATCH 030/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 27b7a2c631..3150732024 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -15,10 +15,9 @@ """Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module.""" from functools import singledispatchmethod -from typing import TYPE_CHECKING, Any +from typing import Any -import xdsl -from xdsl.dialects import builtin, func, scf +from xdsl.dialects import builtin, scf from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import quantum From 2504c1304406ff6536f9ab012865ab40f1e542c8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 10:21:39 -0500 Subject: [PATCH 031/245] clean-up --- .../visualization/construct_circuit_dag.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 3150732024..9ab65459cb 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -57,19 +57,19 @@ def construct(self, module: builtin.ModuleOp) -> None: # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). @visit.register - def visit_operation(self, operation: Operation) -> None: + def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" for region in operation.regions: - self.visit_region(region) + self.visit(region) @visit.register - def visit_region(self, region: Region) -> None: + def _region(self, region: Region) -> None: """Visit an xDSL Region operation.""" for block in region.blocks: - self.visit_block(block) + self.visit(block) @visit.register - def visit_block(self, block: Block) -> None: + def _block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: self.visit(op) From 4e78777df34b15e36583c1585cbecb935fa57928 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 10:58:46 -0500 Subject: [PATCH 032/245] remove unnecessary stuff --- .../visualization/construct_circuit_dag.py | 54 ------------------- 1 file changed, 54 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 9ab65459cb..85b5c85f0c 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -73,57 +73,3 @@ def _block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: self.visit(op) - - # ====================================== - # 3. QUANTUM GATE & STATE PREP HANDLERS - # ====================================== - # Handlers for operations that apply unitary transformations or set-up the quantum state. - - @visit.register - def _unitary_and_state_prep( - self, - op: quantum.CustomOp, - ) -> None: - """Generic handler for unitary gates and quantum state preparation operations.""" - pass - - # ============================================= - # 4. QUANTUM MEASUREMENT HANDLERS - # ============================================= - - @visit.register - def _state_op(self, op: quantum.StateOp) -> None: - """Handler for the terminal state measurement operation.""" - pass - - @visit.register - def _statistical_measurement_ops( - self, - op: quantum.ExpvalOp | quantum.VarianceOp | quantum.ProbsOp | quantum.SampleOp, - ) -> None: - """Handler for statistical measurement operations.""" - pass - - @visit.register - def _projective_measure_op(self, op: quantum.MeasureOp) -> None: - """Handler for the single-qubit projective measurement operation.""" - pass - - # ========================= - # 5. CONTROL FLOW HANDLERS - # ========================= - - @visit.register - def _for_op(self, op: scf.ForOp) -> None: - """Handle an xDSL ForOp operation.""" - pass - - @visit.register - def _while_op(self, op: scf.WhileOp) -> None: - """Handle an xDSL WhileOp operation.""" - pass - - @visit.register - def _if_op(self, op: scf.IfOp) -> None: - """Handle an xDSL IfOp operation.""" - pass From e17c1112364a6056d0cc2a739b998ab69b48f2dc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 11:07:29 -0500 Subject: [PATCH 033/245] add test skeleton --- .../test_construct_circuit_dag.py | 24 +++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index e69de29bb2..336df14c08 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -0,0 +1,24 @@ +# Copyright 2025 Xanadu Quantum Technologies Inc. + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Unit tests for the ConstructCircuitDAG utility.""" + +import pytest + +pytestmark = pytest.mark.usefixtures("requires_xdsl") + +# pylint: disable=wrong-import-position +# This import needs to be after pytest in order to prevent ImportErrors +from catalyst.python_interface.visualization.construct_circuit_dag import ConstructCircuitDAG +from catalyst.python_interface.visualization.dag_builder import DAGBuilder + From 19ca1b5e11ec95a08552fd92aca8be87200d8fff Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 11:28:28 -0500 Subject: [PATCH 034/245] add basic tests --- .../test_construct_circuit_dag.py | 26 +++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 336df14c08..3be7ac78f5 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -13,12 +13,34 @@ # limitations under the License. """Unit tests for the ConstructCircuitDAG utility.""" +from unittest.mock import MagicMock, Mock + import pytest pytestmark = pytest.mark.usefixtures("requires_xdsl") + # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors -from catalyst.python_interface.visualization.construct_circuit_dag import ConstructCircuitDAG -from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from catalyst.python_interface.visualization.construct_circuit_dag import ( + ConstructCircuitDAG, +) +from catalyst.python_interface.visualization.dag_builder import DAGBuilder + + +class TestInitialization: + """Tests that the state is correctly initialized.""" + + def test_dependency_injection(self): + """Tests that relevant dependencies are injected.""" + + mock_dag_builder = Mock(DAGBuilder) + utility = ConstructCircuitDAG(mock_dag_builder) + assert utility.dag_builder is mock_dag_builder + + +class TestRecursiveTraversal: + """Tests that the recursive traversal logic works correctly.""" + def test_entire_module_is_traversed(self): + """Tests that the entire module heirarchy is traversed correctly.""" From 7f88834e3164da84cbbf67d933a2a3777b9a050d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 13:27:05 -0500 Subject: [PATCH 035/245] basic test idea --- .../test_construct_circuit_dag.py | 41 +++++++++++++++++-- 1 file changed, 38 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3be7ac78f5..bad7d9c947 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -13,19 +13,22 @@ # limitations under the License. """Unit tests for the ConstructCircuitDAG utility.""" -from unittest.mock import MagicMock, Mock +from unittest import mock +from unittest.mock import MagicMock, Mock, call import pytest pytestmark = pytest.mark.usefixtures("requires_xdsl") - # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class TestInitialization: @@ -43,4 +46,36 @@ class TestRecursiveTraversal: """Tests that the recursive traversal logic works correctly.""" def test_entire_module_is_traversed(self): - """Tests that the entire module heirarchy is traversed correctly.""" + """Tests that the entire module hierarchy is traversed correctly.""" + + # Create block containing some ops + op = test.TestOp() + block = Block(ops=[op, op]) + # Create region containing some blocks + region = Region(blocks=[block, block]) + # Create op containing the regions + container_op = test.TestOp(regions=[region, region]) + # Create module op to house it all + module_op = ModuleOp(ops=[container_op]) + + mock_dag_builder = Mock(DAGBuilder) + utility = ConstructCircuitDAG(mock_dag_builder) + + # Mock out the visit dispatcher + utility.visit = Mock() + + utility.construct(module_op) + + assert utility.visit.call_count == 5 + + expected_calls = [ + call(container_op), + call(region), + call(region), + call(block), + call(block), + call(op), + call(op), + ] + + utility.visit.assert_has_calls(expected_calls, any_order=False) From 3d48c9291739844e7b0783d2837214e81aed5a8a Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 13:43:07 -0500 Subject: [PATCH 036/245] make visit private --- .../visualization/construct_circuit_dag.py | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 85b5c85f0c..f662e1b7b3 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -41,7 +41,7 @@ def __init__(self, dag_builder: DAGBuilder) -> None: # ================================= @singledispatchmethod - def visit(self, op: Any) -> None: + def _visit(self, op: Any) -> None: """Central dispatch method (Visitor Pattern). Routes the operation 'op' to the specialized handler registered for its type.""" pass @@ -49,27 +49,27 @@ def visit(self, op: Any) -> None: def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module.""" for op in module.ops: - self.visit(op) + self._visit(op) # ======================= # 2. HIERARCHY TRAVERSAL # ======================= # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). - @visit.register + @_visit.register def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" for region in operation.regions: - self.visit(region) + self._visit(region) - @visit.register + @_visit.register def _region(self, region: Region) -> None: """Visit an xDSL Region operation.""" for block in region.blocks: - self.visit(block) + self._visit(block) - @visit.register + @_visit.register def _block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: - self.visit(op) + self._visit(op) From 49b3834a37b18b113149f15d1b95d5994d33a645 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 13:47:34 -0500 Subject: [PATCH 037/245] make visit private --- .../visualization/test_construct_circuit_dag.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index bad7d9c947..88e638ee31 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -62,11 +62,11 @@ def test_entire_module_is_traversed(self): utility = ConstructCircuitDAG(mock_dag_builder) # Mock out the visit dispatcher - utility.visit = Mock() + utility._visit = Mock() utility.construct(module_op) - assert utility.visit.call_count == 5 + assert utility._visit.call_count == 7 expected_calls = [ call(container_op), @@ -78,4 +78,4 @@ def test_entire_module_is_traversed(self): call(op), ] - utility.visit.assert_has_calls(expected_calls, any_order=False) + utility._visit.assert_has_calls(expected_calls, any_order=False) From 80ca59d71df779e98086f4309b22d8b6c50e9011 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 14:12:23 -0500 Subject: [PATCH 038/245] fix tests --- .../test_construct_circuit_dag.py | 23 ++++--------------- 1 file changed, 4 insertions(+), 19 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 88e638ee31..39a74578bb 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -50,32 +50,17 @@ def test_entire_module_is_traversed(self): # Create block containing some ops op = test.TestOp() - block = Block(ops=[op, op]) + block = Block(ops=[op]) # Create region containing some blocks - region = Region(blocks=[block, block]) + region = Region(blocks=[block]) # Create op containing the regions - container_op = test.TestOp(regions=[region, region]) + container_op = test.TestOp(regions=[region]) # Create module op to house it all module_op = ModuleOp(ops=[container_op]) mock_dag_builder = Mock(DAGBuilder) utility = ConstructCircuitDAG(mock_dag_builder) - # Mock out the visit dispatcher - utility._visit = Mock() - utility.construct(module_op) - assert utility._visit.call_count == 7 - - expected_calls = [ - call(container_op), - call(region), - call(region), - call(block), - call(block), - call(op), - call(op), - ] - - utility._visit.assert_has_calls(expected_calls, any_order=False) + # Assert visit was dispatched correct number of times with the correct inputs From 430ceb810fc0015182b624b93d77ba1cc03955b4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 21 Nov 2025 19:37:29 -0500 Subject: [PATCH 039/245] fix tests --- .../test_construct_circuit_dag.py | 42 +++++++++---------- 1 file changed, 20 insertions(+), 22 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 39a74578bb..6d4b7106ea 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -42,25 +42,23 @@ def test_dependency_injection(self): assert utility.dag_builder is mock_dag_builder -class TestRecursiveTraversal: - """Tests that the recursive traversal logic works correctly.""" - - def test_entire_module_is_traversed(self): - """Tests that the entire module hierarchy is traversed correctly.""" - - # Create block containing some ops - op = test.TestOp() - block = Block(ops=[op]) - # Create region containing some blocks - region = Region(blocks=[block]) - # Create op containing the regions - container_op = test.TestOp(regions=[region]) - # Create module op to house it all - module_op = ModuleOp(ops=[container_op]) - - mock_dag_builder = Mock(DAGBuilder) - utility = ConstructCircuitDAG(mock_dag_builder) - - utility.construct(module_op) - - # Assert visit was dispatched correct number of times with the correct inputs +def test_does_not_mutate_module(): + """Test that the module is not mutated.""" + + # Create block containing some ops + op = test.TestOp() + block = Block(ops=[op]) + # Create region containing some blocks + region = Region(blocks=[block]) + # Create op containing the regions + container_op = test.TestOp(regions=[region]) + # Create module op to house it all + module_op = ModuleOp(ops=[container_op]) + + module_op_str_before = str(module_op) + + mock_dag_builder = Mock(DAGBuilder) + utility = ConstructCircuitDAG(mock_dag_builder) + utility.construct(module_op) + + assert str(module_op) == module_op_str_before From 0aff4e5cc57bede900876d6c413dd1f0b7387c2f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 10:08:55 -0500 Subject: [PATCH 040/245] clean-up --- .../visualization/test_construct_circuit_dag.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 6d4b7106ea..5a55cd0ba7 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -13,8 +13,7 @@ # limitations under the License. """Unit tests for the ConstructCircuitDAG utility.""" -from unittest import mock -from unittest.mock import MagicMock, Mock, call +from unittest.mock import Mock import pytest From e22cbd2499f0706bccb639f3112d3e4480693d1e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:08:27 -0500 Subject: [PATCH 041/245] fix: upgrade DAG builders to have get_ methods --- .../visualization/dag_builder.py | 46 +++++++++++++++++-- 1 file changed, 41 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index aeb8b217b6..d59a81b6f4 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -14,7 +14,10 @@ """File that defines the DAGBuilder abstract base class.""" from abc import ABC, abstractmethod -from typing import Any +from typing import Any, TypeAlias + +ClusterID: TypeAlias = str +NodeID: TypeAlias = str class DAGBuilder(ABC): @@ -28,7 +31,11 @@ class DAGBuilder(ABC): @abstractmethod def add_node( - self, node_id: str, node_label: str, parent_graph_id: str | None = None, **node_attrs: Any + self, + node_id: NodeID, + node_label: str, + parent_graph_id: str | None = None, + **node_attrs: Any, ) -> None: """Add a single node to the graph. @@ -42,7 +49,9 @@ def add_node( raise NotImplementedError @abstractmethod - def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + def add_edge( + self, from_node_id: NodeID, to_node_id: NodeID, **edge_attrs: Any + ) -> None: """Add a single directed edge between nodes in the graph. Args: @@ -56,9 +65,9 @@ def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> Non @abstractmethod def add_cluster( self, - cluster_id: str, + cluster_id: ClusterID, node_label: str | None = None, - parent_graph_id: str | None = None, + parent_graph_id: ClusterID | None = None, **cluster_attrs: Any, ) -> None: """Add a single cluster to the graph. @@ -75,6 +84,33 @@ def add_cluster( """ raise NotImplementedError + @abstractmethod + def get_nodes(self) -> dict[NodeID, dict[str, Any]]: + """Retrieve the current set of nodes in the graph. + + Returns: + nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. + """ + raise NotImplementedError + + @abstractmethod + def get_edges(self) -> list[dict[str, Any]]: + """Retrieve the current set of edges in the graph. + + Returns: + edges (list[dict[str, Any]]): A list of edges where each edge contains a dictionary of information for a given edge. + """ + raise NotImplementedError + + @abstractmethod + def get_clusters(self) -> dict[ClusterID, dict[str, Any]]: + """Retrieve the current set of clusters in the graph. + + Returns: + clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. + """ + raise NotImplementedError + @abstractmethod def to_file(self, output_filename: str) -> None: """Save the graph to a file. From acf3da79912c7a0fb5bd493319fb6855eca70ebb Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:12:48 -0500 Subject: [PATCH 042/245] cl --- doc/releases/changelog-dev.md | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index ee7c22c7e3..6e8a6465ff 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -4,6 +4,7 @@ * Compiled programs can be visualized. [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213) + [(#2229)](https://github.com/PennyLaneAI/catalyst/pull/2229) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From bc978adad3e3c51da154549bd24824314860ccfc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:24:25 -0500 Subject: [PATCH 043/245] update pydot to adhere to new base class methods --- .../visualization/pydot_dag_builder.py | 52 ++++++++++++++++++- 1 file changed, 51 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 384537c93e..6f756bf40c 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -57,7 +57,14 @@ def __init__( self._subgraphs: dict[str, pydot.Graph] = {} self._subgraphs["__base__"] = self.graph - _default_attrs: dict = {"fontname": "Helvetica", "penwidth": 2} if attrs is None else attrs + # Internal state for graph structure + self._nodes: dict[str, dict[str, Any]] = {} + self._edges: list[dict[str, Any]] = [] + self._clusters: dict[str, dict[str, Any]] = {} + + _default_attrs: dict = ( + {"fontname": "Helvetica", "penwidth": 2} if attrs is None else attrs + ) self._default_node_attrs: dict = ( { **_default_attrs, @@ -111,6 +118,13 @@ def add_node( self._subgraphs[parent_graph_id].add_node(node) + self._nodes[node_id] = { + "id": node_id, + "label": node_label, + "parent_id": parent_graph_id, + "attrs": dict(node_attrs), + } + def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: """Add a single directed edge between nodes in the graph. @@ -125,6 +139,10 @@ def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> Non edge = pydot.Edge(from_node_id, to_node_id, **edge_attrs) self.graph.add_edge(edge) + self._edges.append( + {"from_id": from_node_id, "to_id": to_node_id, "attrs": dict(edge_attrs)} + ) + def add_cluster( self, cluster_id: str, @@ -177,6 +195,38 @@ def add_cluster( parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id self._subgraphs[parent_graph_id].add_subgraph(cluster) + self._clusters[cluster_id] = { + "id": cluster_id, + "cluster_label": cluster_attrs.get("label"), + "node_label": node_label, + "parent_id": parent_graph_id, + "attrs": dict(cluster_attrs), + } + + def get_nodes(self) -> dict[str, dict[str, Any]]: + """Retrieve the current set of nodes in the graph. + + Returns: + nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. + """ + return self._nodes + + def get_edges(self) -> list[dict[str, Any]]: + """Retrieve the current set of edges in the graph. + + Returns: + edges (list[dict[str, Any]]): A list of edges where each edge contains a dictionary of information for a given edge. + """ + return self._edges + + def get_clusters(self) -> dict[str, dict[str, Any]]: + """Retrieve the current set of clusters in the graph. + + Returns: + clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. + """ + return self._clusters + def to_file(self, output_filename: str) -> None: """Save the graph to a file. From dc4cb1c622ccddba3940e4f3a9e8658dcbebb983 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:37:21 -0500 Subject: [PATCH 044/245] add test skeletons --- .../visualization/pydot_dag_builder.py | 6 +++--- .../visualization/test_dag_builder.py | 15 +++++++++++++++ .../visualization/test_pydot_dag_builder.py | 17 ++++++++++++++--- 3 files changed, 32 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 6f756bf40c..9a81dc55d0 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -112,7 +112,7 @@ def add_node( """ # Use ChainMap so you don't need to construct a new dictionary - node_attrs = ChainMap(node_attrs, self._default_node_attrs) + node_attrs: ChainMap = ChainMap(node_attrs, self._default_node_attrs) node = pydot.Node(node_id, label=node_label, **node_attrs) parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id @@ -135,7 +135,7 @@ def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> Non """ # Use ChainMap so you don't need to construct a new dictionary - edge_attrs = ChainMap(edge_attrs, self._default_edge_attrs) + edge_attrs: ChainMap = ChainMap(edge_attrs, self._default_edge_attrs) edge = pydot.Edge(from_node_id, to_node_id, **edge_attrs) self.graph.add_edge(edge) @@ -163,7 +163,7 @@ def add_cluster( """ # Use ChainMap so you don't need to construct a new dictionary - cluster_attrs = ChainMap(cluster_attrs, self._default_cluster_attrs) + cluster_attrs: ChainMap = ChainMap(cluster_attrs, self._default_cluster_attrs) cluster = pydot.Cluster(graph_name=cluster_id, **cluster_attrs) # Puts the label in a node within the cluster. diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py index 2e935bae1b..7191cf001c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -54,6 +54,15 @@ def add_cluster( ) -> None: return + def get_nodes(self) -> dict[str, dict[str, Any]]: + return {} + + def get_edges(self) -> list[dict[str, Any]]: + return [] + + def get_clusters(self) -> dict[str, dict[str, Any]]: + return {} + def to_file(self, output_filename: str) -> None: return @@ -65,12 +74,18 @@ def to_string(self) -> str: node = dag_builder.add_node("0", "node0") edge = dag_builder.add_edge("0", "1") cluster = dag_builder.add_cluster("0") + nodes = dag_builder.get_nodes() + edges = dag_builder.get_edges() + clusters = dag_builder.get_clusters() render = dag_builder.to_file("test.png") string = dag_builder.to_string() assert node is None + assert nodes == {} assert edge is None + assert edges == [] assert cluster is None + assert clusters == {} assert render is None assert string == "test" diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 11e5e4ac47..426239a929 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -231,9 +231,7 @@ def test_add_cluster_with_attrs(self): assert cluster1.get("penwidth") == 2 assert cluster1.get("fontname") == "Helvetica" - dag_builder.add_cluster( - "1", style="filled", penwidth=10, fillcolor="red" - ) + dag_builder.add_cluster("1", style="filled", penwidth=10, fillcolor="red") cluster2 = dag_builder.graph.get_subgraph("cluster_1")[0] # Make sure we can override @@ -245,6 +243,19 @@ def test_add_cluster_with_attrs(self): assert cluster2.get("fontname") == "Helvetica" +class TestGetMethods: + """Tests the get_* methods.""" + + def test_get_nodes(self): + """Tests that get_nodes works.""" + + def test_get_edges(self): + """Tests that get_edges works.""" + + def test_get_clusters(self): + """Tests that get_clusters works.""" + + class TestOutput: """Test that the graph can be outputted correctly.""" From 7e825d6cb71e5b15ad346c31abee49b87ea88ca3 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:47:47 -0500 Subject: [PATCH 045/245] add tests --- .../visualization/test_pydot_dag_builder.py | 54 +++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 426239a929..81cff2db77 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -248,13 +248,67 @@ class TestGetMethods: def test_get_nodes(self): """Tests that get_nodes works.""" + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + dag_builder.add_cluster("c0", "my_info_node", label="my_cluster") + dag_builder.add_node("1", "node1", parent_graph_id="c0") + + nodes = dag_builder.get_nodes() + assert len(nodes) == 2 + assert len(nodes["0"]) == 4 + assert nodes["0"]["id"] == "0" + assert nodes["0"]["label"] == "node0" + assert nodes["0"]["parent_id"] == "__base__" + assert nodes["0"]["attrs"] == {} + + assert nodes["1"]["id"] == "1" + assert nodes["1"]["label"] == "node1" + assert nodes["1"]["parent_id"] == "c0" + assert nodes["1"]["attrs"] == {} def test_get_edges(self): """Tests that get_edges works.""" + dag_builder = PyDotDAGBuilder() + dag_builder.add_node("0", "node0") + dag_builder.add_node("1", "node1") + dag_builder.add_edge("0", "1") + + edges = dag_builder.get_edges() + assert len(edges) == 1 + assert edges[0]["from_id"] == "0" + assert edges[0]["to_id"] == "0" + assert edges[0]["attrs"] == {} + def test_get_clusters(self): """Tests that get_clusters works.""" + dag_builder = PyDotDAGBuilder() + dag_builder.add_cluster("0", "my_info_node", label="my_cluster") + + clusters = dag_builder.get_clusters() + assert len(clusters) == 1 + assert len(clusters["0"]) == 5 + assert clusters["0"]["id"] == "0" + assert clusters["0"]["cluster_label"] == "my_cluster" + assert clusters["0"]["node_label"] == "my_info_node" + assert clusters["0"]["parent_id"] == "__base__" + assert clusters["0"]["attrs"] == {} + + dag_builder.add_cluster( + "1", "my_other_info_node", parent_graph_id="0", label="my_nested_cluster" + ) + + clusters = dag_builder.get_clusters() + assert len(clusters) == 2 + assert len(clusters["1"]) == 5 + assert clusters["1"]["id"] == "0" + assert clusters["1"]["cluster_label"] == "my_nested_cluster" + assert clusters["1"]["node_label"] == "my_other_info_node" + assert clusters["1"]["parent_id"] == "0" + assert clusters["1"]["attrs"] == {} + class TestOutput: """Test that the graph can be outputted correctly.""" From 41fc4d183e6699c1520150826e58cb0f2f8c7345 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Mon, 24 Nov 2025 13:48:14 -0500 Subject: [PATCH 046/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 81cff2db77..d06b548478 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -278,7 +278,7 @@ def test_get_edges(self): edges = dag_builder.get_edges() assert len(edges) == 1 assert edges[0]["from_id"] == "0" - assert edges[0]["to_id"] == "0" + assert edges[0]["to_id"] == "1" assert edges[0]["attrs"] == {} def test_get_clusters(self): From 1868346de42eac6ba9bb8e5956df92bce67edf29 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:49:07 -0500 Subject: [PATCH 047/245] add tests --- .../visualization/test_pydot_dag_builder.py | 18 +++++++++++------- 1 file changed, 11 insertions(+), 7 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index d06b548478..805bacb6ab 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -255,8 +255,10 @@ def test_get_nodes(self): dag_builder.add_node("1", "node1", parent_graph_id="c0") nodes = dag_builder.get_nodes() + assert len(nodes) == 2 assert len(nodes["0"]) == 4 + assert nodes["0"]["id"] == "0" assert nodes["0"]["label"] == "node0" assert nodes["0"]["parent_id"] == "__base__" @@ -276,7 +278,9 @@ def test_get_edges(self): dag_builder.add_edge("0", "1") edges = dag_builder.get_edges() + assert len(edges) == 1 + assert edges[0]["from_id"] == "0" assert edges[0]["to_id"] == "1" assert edges[0]["attrs"] == {} @@ -288,7 +292,13 @@ def test_get_clusters(self): dag_builder.add_cluster("0", "my_info_node", label="my_cluster") clusters = dag_builder.get_clusters() - assert len(clusters) == 1 + + dag_builder.add_cluster( + "1", "my_other_info_node", parent_graph_id="0", label="my_nested_cluster" + ) + clusters = dag_builder.get_clusters() + assert len(clusters) == 2 + assert len(clusters["0"]) == 5 assert clusters["0"]["id"] == "0" assert clusters["0"]["cluster_label"] == "my_cluster" @@ -296,12 +306,6 @@ def test_get_clusters(self): assert clusters["0"]["parent_id"] == "__base__" assert clusters["0"]["attrs"] == {} - dag_builder.add_cluster( - "1", "my_other_info_node", parent_graph_id="0", label="my_nested_cluster" - ) - - clusters = dag_builder.get_clusters() - assert len(clusters) == 2 assert len(clusters["1"]) == 5 assert clusters["1"]["id"] == "0" assert clusters["1"]["cluster_label"] == "my_nested_cluster" From fd8d72116276b3e56e9b8f60078ddd2117845391 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:50:00 -0500 Subject: [PATCH 048/245] fix tests --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 805bacb6ab..28b04f5f41 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -251,7 +251,7 @@ def test_get_nodes(self): dag_builder = PyDotDAGBuilder() dag_builder.add_node("0", "node0") - dag_builder.add_cluster("c0", "my_info_node", label="my_cluster") + dag_builder.add_cluster("c0") dag_builder.add_node("1", "node1", parent_graph_id="c0") nodes = dag_builder.get_nodes() From 284ba07eab92d8e6e14d16ac00259f78d23a2909 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 13:59:50 -0500 Subject: [PATCH 049/245] update tests --- .../visualization/test_pydot_dag_builder.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 28b04f5f41..2730efa412 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -250,7 +250,7 @@ def test_get_nodes(self): """Tests that get_nodes works.""" dag_builder = PyDotDAGBuilder() - dag_builder.add_node("0", "node0") + dag_builder.add_node("0", "node0", fillcolor="red") dag_builder.add_cluster("c0") dag_builder.add_node("1", "node1", parent_graph_id="c0") @@ -262,7 +262,7 @@ def test_get_nodes(self): assert nodes["0"]["id"] == "0" assert nodes["0"]["label"] == "node0" assert nodes["0"]["parent_id"] == "__base__" - assert nodes["0"]["attrs"] == {} + assert nodes["0"]["attrs"] == {"fillcolor": "red"} assert nodes["1"]["id"] == "1" assert nodes["1"]["label"] == "node1" @@ -275,7 +275,7 @@ def test_get_edges(self): dag_builder = PyDotDAGBuilder() dag_builder.add_node("0", "node0") dag_builder.add_node("1", "node1") - dag_builder.add_edge("0", "1") + dag_builder.add_edge("0", "1", penwidth=10) edges = dag_builder.get_edges() @@ -283,13 +283,13 @@ def test_get_edges(self): assert edges[0]["from_id"] == "0" assert edges[0]["to_id"] == "1" - assert edges[0]["attrs"] == {} + assert edges[0]["attrs"] == {"penwidth": 10} def test_get_clusters(self): """Tests that get_clusters works.""" dag_builder = PyDotDAGBuilder() - dag_builder.add_cluster("0", "my_info_node", label="my_cluster") + dag_builder.add_cluster("0", "my_info_node", label="my_cluster", penwidth=10) clusters = dag_builder.get_clusters() @@ -304,7 +304,7 @@ def test_get_clusters(self): assert clusters["0"]["cluster_label"] == "my_cluster" assert clusters["0"]["node_label"] == "my_info_node" assert clusters["0"]["parent_id"] == "__base__" - assert clusters["0"]["attrs"] == {} + assert clusters["0"]["attrs"] == {"penwidth": 10} assert len(clusters["1"]) == 5 assert clusters["1"]["id"] == "0" From 95ace870e83f989a1c1e24aafde9465ce22b00b0 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 14:02:05 -0500 Subject: [PATCH 050/245] fix documentation --- .../catalyst/python_interface/visualization/dag_builder.py | 3 ++- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 2 files changed, 3 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index d59a81b6f4..492976ed5f 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -78,7 +78,8 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. node_label (str | None): The text to display on an information node within the cluster when rendered. - parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. + parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be + placed on the base graph. **cluster_attrs (Any): Any additional styling keyword arguments. """ diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 9a81dc55d0..7073f9a52e 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -158,7 +158,7 @@ def add_cluster( Args: cluster_id (str): Unique cluster ID to identify this cluster. node_label (str | None): The text to display on the information node within the cluster when rendered. - parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. + parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be positioned on the base graph. **cluster_attrs (Any): Any additional styling keyword arguments. """ From ec9caaeddb1b10d07dc8d3843a2a91f25b04d6e4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 14:04:18 -0500 Subject: [PATCH 051/245] fix documentation --- frontend/catalyst/python_interface/visualization/dag_builder.py | 2 +- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 492976ed5f..db118b4dfb 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -99,7 +99,7 @@ def get_edges(self) -> list[dict[str, Any]]: """Retrieve the current set of edges in the graph. Returns: - edges (list[dict[str, Any]]): A list of edges where each edge contains a dictionary of information for a given edge. + edges (list[dict[str, Any]]): A list of edges where each element in the list contains a dictionary of edge information. """ raise NotImplementedError diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 7073f9a52e..2e0ba196a4 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -215,7 +215,7 @@ def get_edges(self) -> list[dict[str, Any]]: """Retrieve the current set of edges in the graph. Returns: - edges (list[dict[str, Any]]): A list of edges where each edge contains a dictionary of information for a given edge. + edges (list[dict[str, Any]]): A list of edges where each element in the list contains a dictionary of edge information. """ return self._edges From 99bd602308c03d44d763affd9025de75b8b5d62d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 14:19:07 -0500 Subject: [PATCH 052/245] fix tests --- .../visualization/test_pydot_dag_builder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 2730efa412..cff4f72207 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -262,7 +262,7 @@ def test_get_nodes(self): assert nodes["0"]["id"] == "0" assert nodes["0"]["label"] == "node0" assert nodes["0"]["parent_id"] == "__base__" - assert nodes["0"]["attrs"] == {"fillcolor": "red"} + assert nodes["0"]["attrs"]["fillcolor"] == "red" assert nodes["1"]["id"] == "1" assert nodes["1"]["label"] == "node1" @@ -283,7 +283,7 @@ def test_get_edges(self): assert edges[0]["from_id"] == "0" assert edges[0]["to_id"] == "1" - assert edges[0]["attrs"] == {"penwidth": 10} + assert edges[0]["attrs"]["penwidth"] == 10 def test_get_clusters(self): """Tests that get_clusters works.""" @@ -304,7 +304,7 @@ def test_get_clusters(self): assert clusters["0"]["cluster_label"] == "my_cluster" assert clusters["0"]["node_label"] == "my_info_node" assert clusters["0"]["parent_id"] == "__base__" - assert clusters["0"]["attrs"] == {"penwidth": 10} + assert clusters["0"]["attrs"]["penwidth"] == 10 assert len(clusters["1"]) == 5 assert clusters["1"]["id"] == "0" From 7622b502560dd67f282162f1b4539ef19e568985 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 14:35:24 -0500 Subject: [PATCH 053/245] whoops --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index cff4f72207..dbec2ced99 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -267,7 +267,6 @@ def test_get_nodes(self): assert nodes["1"]["id"] == "1" assert nodes["1"]["label"] == "node1" assert nodes["1"]["parent_id"] == "c0" - assert nodes["1"]["attrs"] == {} def test_get_edges(self): """Tests that get_edges works.""" @@ -311,7 +310,6 @@ def test_get_clusters(self): assert clusters["1"]["cluster_label"] == "my_nested_cluster" assert clusters["1"]["node_label"] == "my_other_info_node" assert clusters["1"]["parent_id"] == "0" - assert clusters["1"]["attrs"] == {} class TestOutput: From 86e8389feaa02b555ec1a1498558a46d4a17a8f4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 14:58:15 -0500 Subject: [PATCH 054/245] whoops --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index dbec2ced99..40fce1d641 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -306,7 +306,7 @@ def test_get_clusters(self): assert clusters["0"]["attrs"]["penwidth"] == 10 assert len(clusters["1"]) == 5 - assert clusters["1"]["id"] == "0" + assert clusters["1"]["id"] == "1" assert clusters["1"]["cluster_label"] == "my_nested_cluster" assert clusters["1"]["node_label"] == "my_other_info_node" assert clusters["1"]["parent_id"] == "0" From a5dc1ac0f51812b1d40d06c4fa3185d329499574 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 15:12:55 -0500 Subject: [PATCH 055/245] rename a bunch of stuff --- .../visualization/dag_builder.py | 28 +++++----- .../visualization/pydot_dag_builder.py | 52 +++++++++---------- .../visualization/test_dag_builder.py | 14 +++-- .../visualization/test_pydot_dag_builder.py | 12 ++--- 4 files changed, 51 insertions(+), 55 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index db118b4dfb..6ea69de6e7 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -32,31 +32,29 @@ class DAGBuilder(ABC): @abstractmethod def add_node( self, - node_id: NodeID, - node_label: str, - parent_graph_id: str | None = None, + id: NodeID, + label: str, + cluster_id: ClusterID | None = None, **node_attrs: Any, ) -> None: """Add a single node to the graph. Args: - node_id (str): Unique node ID to identify this node. - node_label (str): The text to display on the node when rendered. - parent_graph_id (str | None): Optional ID of the cluster this node belongs to. + id (str): Unique node ID to identify this node. + label (str): The text to display on the node when rendered. + cluster_id (str | None): Optional ID of the cluster this node belongs to. **node_attrs (Any): Any additional styling keyword arguments. """ raise NotImplementedError @abstractmethod - def add_edge( - self, from_node_id: NodeID, to_node_id: NodeID, **edge_attrs: Any - ) -> None: + def add_edge(self, from_id: NodeID, to_id: NodeID, **edge_attrs: Any) -> None: """Add a single directed edge between nodes in the graph. Args: - from_node_id (str): The unique ID of the source node. - to_node_id (str): The unique ID of the destination node. + from_id (str): The unique ID of the source node. + to_id (str): The unique ID of the destination node. **edge_attrs (Any): Any additional styling keyword arguments. """ @@ -65,9 +63,9 @@ def add_edge( @abstractmethod def add_cluster( self, - cluster_id: ClusterID, + id: ClusterID, node_label: str | None = None, - parent_graph_id: ClusterID | None = None, + cluster_id: ClusterID | None = None, **cluster_attrs: Any, ) -> None: """Add a single cluster to the graph. @@ -76,9 +74,9 @@ def add_cluster( within it are visually and logically grouped. Args: - cluster_id (str): Unique cluster ID to identify this cluster. + id (str): Unique cluster ID to identify this cluster. node_label (str | None): The text to display on an information node within the cluster when rendered. - parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be + cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be placed on the base graph. **cluster_attrs (Any): Any additional styling keyword arguments. diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 2e0ba196a4..57b488b530 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -97,57 +97,57 @@ def __init__( def add_node( self, - node_id: str, - node_label: str, - parent_graph_id: str | None = None, + id: str, + label: str, + cluster_id: str | None = None, **node_attrs: Any, ) -> None: """Add a single node to the graph. Args: - node_id (str): Unique node ID to identify this node. - node_label (str): The text to display on the node when rendered. - parent_graph_id (str | None): Optional ID of the cluster this node belongs to. + id (str): Unique node ID to identify this node. + label (str): The text to display on the node when rendered. + cluster_id (str | None): Optional ID of the cluster this node belongs to. **node_attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary node_attrs: ChainMap = ChainMap(node_attrs, self._default_node_attrs) - node = pydot.Node(node_id, label=node_label, **node_attrs) - parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id + node = pydot.Node(id, label=label, **node_attrs) + parent_graph_id = "__base__" if cluster_id is None else cluster_id self._subgraphs[parent_graph_id].add_node(node) - self._nodes[node_id] = { - "id": node_id, - "label": node_label, - "parent_id": parent_graph_id, + self._nodes[id] = { + "id": id, + "label": label, + "cluster_id": cluster_id, "attrs": dict(node_attrs), } - def add_edge(self, from_node_id: str, to_node_id: str, **edge_attrs: Any) -> None: + def add_edge(self, from_id: str, to_id: str, **edge_attrs: Any) -> None: """Add a single directed edge between nodes in the graph. Args: - from_node_id (str): The unique ID of the source node. - to_node_id (str): The unique ID of the destination node. + from_id (str): The unique ID of the source node. + to_id (str): The unique ID of the destination node. **edge_attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary edge_attrs: ChainMap = ChainMap(edge_attrs, self._default_edge_attrs) - edge = pydot.Edge(from_node_id, to_node_id, **edge_attrs) + edge = pydot.Edge(from_id, to_id, **edge_attrs) self.graph.add_edge(edge) self._edges.append( - {"from_id": from_node_id, "to_id": to_node_id, "attrs": dict(edge_attrs)} + {"from_id": from_id, "to_id": to_id, "attrs": dict(edge_attrs)} ) def add_cluster( self, - cluster_id: str, + id: str, node_label: str | None = None, - parent_graph_id: str | None = None, + cluster_id: str | None = None, **cluster_attrs: Any, ) -> None: """Add a single cluster to the graph. @@ -156,15 +156,15 @@ def add_cluster( within it are visually and logically grouped. Args: - cluster_id (str): Unique cluster ID to identify this cluster. + id (str): Unique cluster ID to identify this cluster. node_label (str | None): The text to display on the information node within the cluster when rendered. - parent_graph_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be positioned on the base graph. + cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be positioned on the base graph. **cluster_attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary cluster_attrs: ChainMap = ChainMap(cluster_attrs, self._default_cluster_attrs) - cluster = pydot.Cluster(graph_name=cluster_id, **cluster_attrs) + cluster = pydot.Cluster(graph_name=id, **cluster_attrs) # Puts the label in a node within the cluster. # Ensures that any edges connecting nodes through the cluster @@ -190,13 +190,13 @@ def add_cluster( cluster.add_subgraph(rank_subgraph) cluster.add_node(node) - self._subgraphs[cluster_id] = cluster + self._subgraphs[id] = cluster - parent_graph_id = "__base__" if parent_graph_id is None else parent_graph_id + parent_graph_id = "__base__" if cluster_id is None else cluster_id self._subgraphs[parent_graph_id].add_subgraph(cluster) - self._clusters[cluster_id] = { - "id": cluster_id, + self._clusters[id] = { + "id": id, "cluster_label": cluster_attrs.get("label"), "node_label": node_label, "parent_id": parent_graph_id, diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py index 7191cf001c..3b7d61c49d 100644 --- a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -33,23 +33,21 @@ class ConcreteDAGBuilder(DAGBuilder): def add_node( self, - node_id: str, - node_label: str, - parent_graph_id: str | None = None, + id: str, + label: str, + cluster_id: str | None = None, **node_attrs: Any, ) -> None: return - def add_edge( - self, from_node_id: str, to_node_id: str, **edge_attrs: Any - ) -> None: + def add_edge(self, from_id: str, to_id: str, **edge_attrs: Any) -> None: return def add_cluster( self, - cluster_id: str, + id: str, node_label: str | None = None, - parent_graph_id: str | None = None, + cluster_id: str | None = None, **cluster_attrs: Any, ) -> None: return diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 40fce1d641..b0a17441ea 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -90,7 +90,7 @@ def test_add_node_to_parent_graph(self): dag_builder.add_cluster("c0") # Create node inside cluster - dag_builder.add_node("1", "node1", parent_graph_id="c0") + dag_builder.add_node("1", "node1", cluster_id="c0") # Verify graph structure root_graph = dag_builder.graph @@ -120,11 +120,11 @@ def test_add_cluster_to_parent_graph(self): dag_builder.add_cluster("c0") # Level 1 (Inside c0): Add node on outer cluster and create new cluster on top - dag_builder.add_node("n_outer", "node_outer", parent_graph_id="c0") - dag_builder.add_cluster("c1", parent_graph_id="c0") + dag_builder.add_node("n_outer", "node_outer", cluster_id="c0") + dag_builder.add_cluster("c1", cluster_id="c0") # Level 2 (Inside c1): Add node on second cluster - dag_builder.add_node("n_inner", "node_inner", parent_graph_id="c1") + dag_builder.add_node("n_inner", "node_inner", cluster_id="c1") root_graph = dag_builder.graph @@ -252,7 +252,7 @@ def test_get_nodes(self): dag_builder.add_node("0", "node0", fillcolor="red") dag_builder.add_cluster("c0") - dag_builder.add_node("1", "node1", parent_graph_id="c0") + dag_builder.add_node("1", "node1", cluster_id="c0") nodes = dag_builder.get_nodes() @@ -293,7 +293,7 @@ def test_get_clusters(self): clusters = dag_builder.get_clusters() dag_builder.add_cluster( - "1", "my_other_info_node", parent_graph_id="0", label="my_nested_cluster" + "1", "my_other_info_node", cluster_id="0", label="my_nested_cluster" ) clusters = dag_builder.get_clusters() assert len(clusters) == 2 From fe7ef4772df9f9337b776cf46998e7cf32b2e7be Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 15:13:51 -0500 Subject: [PATCH 056/245] fix documentation --- .../catalyst/python_interface/visualization/dag_builder.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 6ea69de6e7..becd926b91 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -42,7 +42,8 @@ def add_node( Args: id (str): Unique node ID to identify this node. label (str): The text to display on the node when rendered. - cluster_id (str | None): Optional ID of the cluster this node belongs to. + cluster_id (str | None): Optional ID of the cluster this node belongs to. If `None`, this node gets + added on the base graph. **node_attrs (Any): Any additional styling keyword arguments. """ From b2a13cfc8fd64fc9ab94f2be3005407a6617a8a1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 15:14:49 -0500 Subject: [PATCH 057/245] rename a bunch of stuff --- .../visualization/pydot_dag_builder.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 57b488b530..831b181fec 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -114,9 +114,9 @@ def add_node( # Use ChainMap so you don't need to construct a new dictionary node_attrs: ChainMap = ChainMap(node_attrs, self._default_node_attrs) node = pydot.Node(id, label=label, **node_attrs) - parent_graph_id = "__base__" if cluster_id is None else cluster_id + cluster_id = "__base__" if cluster_id is None else cluster_id - self._subgraphs[parent_graph_id].add_node(node) + self._subgraphs[cluster_id].add_node(node) self._nodes[id] = { "id": id, @@ -192,14 +192,14 @@ def add_cluster( self._subgraphs[id] = cluster - parent_graph_id = "__base__" if cluster_id is None else cluster_id - self._subgraphs[parent_graph_id].add_subgraph(cluster) + cluster_id = "__base__" if cluster_id is None else cluster_id + self._subgraphs[cluster_id].add_subgraph(cluster) self._clusters[id] = { "id": id, "cluster_label": cluster_attrs.get("label"), "node_label": node_label, - "parent_id": parent_graph_id, + "parent_id": cluster_id, "attrs": dict(cluster_attrs), } From abfd9328fb046846ee9de01332d8d3ad6258fefb Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 15:19:09 -0500 Subject: [PATCH 058/245] add dev comment --- .../catalyst/python_interface/visualization/pydot_dag_builder.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 831b181fec..4cd63de03f 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -54,6 +54,7 @@ def __init__( graph_type="digraph", rankdir="TB", compound="true", strict=True ) # Create cache for easy look-up + # TODO: Get rid of this and use self._clusters if possible self._subgraphs: dict[str, pydot.Graph] = {} self._subgraphs["__base__"] = self.graph From a481a048748c4b4ed0e27a1b584a593cc3f44a4d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 15:22:03 -0500 Subject: [PATCH 059/245] rename --- .../visualization/dag_builder.py | 12 ++++++------ .../visualization/pydot_dag_builder.py | 18 +++++++++--------- .../visualization/test_dag_builder.py | 6 +++--- .../visualization/test_pydot_dag_builder.py | 6 +++--- 4 files changed, 21 insertions(+), 21 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index becd926b91..8a2fff6960 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -35,7 +35,7 @@ def add_node( id: NodeID, label: str, cluster_id: ClusterID | None = None, - **node_attrs: Any, + **attrs: Any, ) -> None: """Add a single node to the graph. @@ -44,19 +44,19 @@ def add_node( label (str): The text to display on the node when rendered. cluster_id (str | None): Optional ID of the cluster this node belongs to. If `None`, this node gets added on the base graph. - **node_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ raise NotImplementedError @abstractmethod - def add_edge(self, from_id: NodeID, to_id: NodeID, **edge_attrs: Any) -> None: + def add_edge(self, from_id: NodeID, to_id: NodeID, **attrs: Any) -> None: """Add a single directed edge between nodes in the graph. Args: from_id (str): The unique ID of the source node. to_id (str): The unique ID of the destination node. - **edge_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ raise NotImplementedError @@ -67,7 +67,7 @@ def add_cluster( id: ClusterID, node_label: str | None = None, cluster_id: ClusterID | None = None, - **cluster_attrs: Any, + **attrs: Any, ) -> None: """Add a single cluster to the graph. @@ -79,7 +79,7 @@ def add_cluster( node_label (str | None): The text to display on an information node within the cluster when rendered. cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be placed on the base graph. - **cluster_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ raise NotImplementedError diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 4cd63de03f..47e9f584f9 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -101,7 +101,7 @@ def add_node( id: str, label: str, cluster_id: str | None = None, - **node_attrs: Any, + **attrs: Any, ) -> None: """Add a single node to the graph. @@ -109,11 +109,11 @@ def add_node( id (str): Unique node ID to identify this node. label (str): The text to display on the node when rendered. cluster_id (str | None): Optional ID of the cluster this node belongs to. - **node_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary - node_attrs: ChainMap = ChainMap(node_attrs, self._default_node_attrs) + node_attrs: ChainMap = ChainMap(attrs, self._default_node_attrs) node = pydot.Node(id, label=label, **node_attrs) cluster_id = "__base__" if cluster_id is None else cluster_id @@ -126,17 +126,17 @@ def add_node( "attrs": dict(node_attrs), } - def add_edge(self, from_id: str, to_id: str, **edge_attrs: Any) -> None: + def add_edge(self, from_id: str, to_id: str, **attrs: Any) -> None: """Add a single directed edge between nodes in the graph. Args: from_id (str): The unique ID of the source node. to_id (str): The unique ID of the destination node. - **edge_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary - edge_attrs: ChainMap = ChainMap(edge_attrs, self._default_edge_attrs) + edge_attrs: ChainMap = ChainMap(attrs, self._default_edge_attrs) edge = pydot.Edge(from_id, to_id, **edge_attrs) self.graph.add_edge(edge) @@ -149,7 +149,7 @@ def add_cluster( id: str, node_label: str | None = None, cluster_id: str | None = None, - **cluster_attrs: Any, + **attrs: Any, ) -> None: """Add a single cluster to the graph. @@ -160,11 +160,11 @@ def add_cluster( id (str): Unique cluster ID to identify this cluster. node_label (str | None): The text to display on the information node within the cluster when rendered. cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be positioned on the base graph. - **cluster_attrs (Any): Any additional styling keyword arguments. + **attrs (Any): Any additional styling keyword arguments. """ # Use ChainMap so you don't need to construct a new dictionary - cluster_attrs: ChainMap = ChainMap(cluster_attrs, self._default_cluster_attrs) + cluster_attrs: ChainMap = ChainMap(attrs, self._default_cluster_attrs) cluster = pydot.Cluster(graph_name=id, **cluster_attrs) # Puts the label in a node within the cluster. diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py index 3b7d61c49d..2ff5b8255f 100644 --- a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -36,11 +36,11 @@ def add_node( id: str, label: str, cluster_id: str | None = None, - **node_attrs: Any, + **attrs: Any, ) -> None: return - def add_edge(self, from_id: str, to_id: str, **edge_attrs: Any) -> None: + def add_edge(self, from_id: str, to_id: str, **attrs: Any) -> None: return def add_cluster( @@ -48,7 +48,7 @@ def add_cluster( id: str, node_label: str | None = None, cluster_id: str | None = None, - **cluster_attrs: Any, + **attrs: Any, ) -> None: return diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index b0a17441ea..808dcedc87 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -176,7 +176,7 @@ def test_default_graph_attrs(self): def test_add_node_with_attrs(self): """Tests that default attributes are applied and can be overridden.""" dag_builder = PyDotDAGBuilder( - node_attrs={"fillcolor": "lightblue", "penwidth": 3} + attrs={"fillcolor": "lightblue", "penwidth": 3} ) # Defaults @@ -194,7 +194,7 @@ def test_add_node_with_attrs(self): @pytest.mark.unit def test_add_edge_with_attrs(self): """Tests that default attributes are applied and can be overridden.""" - dag_builder = PyDotDAGBuilder(edge_attrs={"color": "lightblue4", "penwidth": 3}) + dag_builder = PyDotDAGBuilder(attrs={"color": "lightblue4", "penwidth": 3}) dag_builder.add_node("0", "node0") dag_builder.add_node("1", "node1") @@ -214,7 +214,7 @@ def test_add_edge_with_attrs(self): def test_add_cluster_with_attrs(self): """Tests that default cluster attributes are applied and can be overridden.""" dag_builder = PyDotDAGBuilder( - cluster_attrs={ + attrs={ "style": "solid", "fillcolor": None, "penwidth": 2, From 9a18e5c25bb408e41828e3d28878df8831ba0a86 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 16:10:16 -0500 Subject: [PATCH 060/245] update test --- .../python_interface/visualization/test_pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 808dcedc87..12c44df287 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -261,12 +261,12 @@ def test_get_nodes(self): assert nodes["0"]["id"] == "0" assert nodes["0"]["label"] == "node0" - assert nodes["0"]["parent_id"] == "__base__" + assert nodes["0"]["cluster_id"] == "__base__" assert nodes["0"]["attrs"]["fillcolor"] == "red" assert nodes["1"]["id"] == "1" assert nodes["1"]["label"] == "node1" - assert nodes["1"]["parent_id"] == "c0" + assert nodes["1"]["cluster_id"] == "c0" def test_get_edges(self): """Tests that get_edges works.""" From 0f6ab76e63e5573df041ffaa80605c374e9e07c6 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 24 Nov 2025 16:56:56 -0500 Subject: [PATCH 061/245] add immutability tests --- .../visualization/pydot_dag_builder.py | 6 +- .../visualization/test_pydot_dag_builder.py | 59 ++++++++++++++++++- 2 files changed, 59 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 47e9f584f9..ea6417d2d4 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -210,7 +210,7 @@ def get_nodes(self) -> dict[str, dict[str, Any]]: Returns: nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. """ - return self._nodes + return self._nodes.copy() def get_edges(self) -> list[dict[str, Any]]: """Retrieve the current set of edges in the graph. @@ -218,7 +218,7 @@ def get_edges(self) -> list[dict[str, Any]]: Returns: edges (list[dict[str, Any]]): A list of edges where each element in the list contains a dictionary of edge information. """ - return self._edges + return self._edges.copy() def get_clusters(self) -> dict[str, dict[str, Any]]: """Retrieve the current set of clusters in the graph. @@ -226,7 +226,7 @@ def get_clusters(self) -> dict[str, dict[str, Any]]: Returns: clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. """ - return self._clusters + return self._clusters.copy() def to_file(self, output_filename: str) -> None: """Save the graph to a file. diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 12c44df287..db6e9fb4da 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -17,6 +17,8 @@ import pytest +from frontend.catalyst.python_interface.visualization import dag_builder + pydot = pytest.importorskip("pydot") pytestmark = pytest.mark.usefixtures("requires_xdsl") # pylint: disable=wrong-import-position @@ -175,9 +177,7 @@ def test_default_graph_attrs(self): @pytest.mark.unit def test_add_node_with_attrs(self): """Tests that default attributes are applied and can be overridden.""" - dag_builder = PyDotDAGBuilder( - attrs={"fillcolor": "lightblue", "penwidth": 3} - ) + dag_builder = PyDotDAGBuilder(attrs={"fillcolor": "lightblue", "penwidth": 3}) # Defaults dag_builder.add_node("0", "node0") @@ -268,6 +268,23 @@ def test_get_nodes(self): assert nodes["1"]["label"] == "node1" assert nodes["1"]["cluster_id"] == "c0" + def test_get_nodes_doesnt_mutate(self): + """Tests that get_nodes doesn't mutate state""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + + old_nodes = dag_builder.get_nodes() + + dag_builder.add_node("1", "node1") + + new_nodes = dag_builder.get_nodes() + + assert old_nodes is not new_nodes + assert len(old_nodes) == 1 + assert len(new_nodes) == 2 + def test_get_edges(self): """Tests that get_edges works.""" @@ -284,6 +301,25 @@ def test_get_edges(self): assert edges[0]["to_id"] == "1" assert edges[0]["attrs"]["penwidth"] == 10 + def test_get_edges_doesnt_mutate(self): + """Tests that get_edges doesn't mutated.""" + + dag_builder = PyDotDAGBuilder() + dag_builder.add_node("0", "node0") + dag_builder.add_node("1", "node1") + dag_builder.add_edge("0", "1") + + old_edges = dag_builder.get_edges() + + dag_builder.add_node("2", "node2") + dag_builder.add_edge("1", "2") + + new_edges = dag_builder.get_edges() + + assert old_edges is not new_edges + assert len(old_edges) == 1 + assert len(new_edges) == 2 + def test_get_clusters(self): """Tests that get_clusters works.""" @@ -311,6 +347,23 @@ def test_get_clusters(self): assert clusters["1"]["node_label"] == "my_other_info_node" assert clusters["1"]["parent_id"] == "0" + def test_get_clusters_doesnt_mutate(self): + """Tests that get_clusters doesn't mutate state""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_cluster("0") + + old_clusters = dag_builder.get_clusters() + + dag_builder.add_cluster("1") + + new_clusters = dag_builder.get_clusters() + + assert old_clusters is not new_clusters + assert len(old_clusters) == 1 + assert len(new_clusters) == 2 + class TestOutput: """Test that the graph can be outputted correctly.""" From 6d15d7be495875fe325e7b0d7b010ed97ca19e7f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 08:54:05 -0500 Subject: [PATCH 062/245] clean-up --- .../python_interface/visualization/construct_circuit_dag.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index f662e1b7b3..4ab442e8cb 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -52,9 +52,8 @@ def construct(self, module: builtin.ModuleOp) -> None: self._visit(op) # ======================= - # 2. HIERARCHY TRAVERSAL + # 2. IR TRAVERSAL # ======================= - # These methods navigate the recursive IR hierarchy (Op -> Region -> Block -> Op). @_visit.register def _operation(self, operation: Operation) -> None: From 888d025dafdb04b2d74314c0d5d00aef61e59be9 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 08:55:58 -0500 Subject: [PATCH 063/245] clean-up --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 4ab442e8cb..419943c361 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -44,7 +44,6 @@ def __init__(self, dag_builder: DAGBuilder) -> None: def _visit(self, op: Any) -> None: """Central dispatch method (Visitor Pattern). Routes the operation 'op' to the specialized handler registered for its type.""" - pass def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module.""" From d4b34f96d5995c8252faa63c61d84745a6597461 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:01:28 -0500 Subject: [PATCH 064/245] whoops --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index ea6417d2d4..e8abda83da 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -200,7 +200,7 @@ def add_cluster( "id": id, "cluster_label": cluster_attrs.get("label"), "node_label": node_label, - "parent_id": cluster_id, + "cluster_id": cluster_id, "attrs": dict(cluster_attrs), } From 33aa334ba1286f8f55b6980a7951bbc8a0747bae Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:15:51 -0500 Subject: [PATCH 065/245] fix test --- .../python_interface/visualization/test_pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index db6e9fb4da..761be2da21 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -338,14 +338,14 @@ def test_get_clusters(self): assert clusters["0"]["id"] == "0" assert clusters["0"]["cluster_label"] == "my_cluster" assert clusters["0"]["node_label"] == "my_info_node" - assert clusters["0"]["parent_id"] == "__base__" + assert clusters["0"]["cluster_id"] == "__base__" assert clusters["0"]["attrs"]["penwidth"] == 10 assert len(clusters["1"]) == 5 assert clusters["1"]["id"] == "1" assert clusters["1"]["cluster_label"] == "my_nested_cluster" assert clusters["1"]["node_label"] == "my_other_info_node" - assert clusters["1"]["parent_id"] == "0" + assert clusters["1"]["cluster_id"] == "0" def test_get_clusters_doesnt_mutate(self): """Tests that get_clusters doesn't mutate state""" From d50cfdc75bd2fe6fdbe851bd4927b45b846204b5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:16:58 -0500 Subject: [PATCH 066/245] whoops --- .../python_interface/visualization/test_pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index db6e9fb4da..761be2da21 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -338,14 +338,14 @@ def test_get_clusters(self): assert clusters["0"]["id"] == "0" assert clusters["0"]["cluster_label"] == "my_cluster" assert clusters["0"]["node_label"] == "my_info_node" - assert clusters["0"]["parent_id"] == "__base__" + assert clusters["0"]["cluster_id"] == "__base__" assert clusters["0"]["attrs"]["penwidth"] == 10 assert len(clusters["1"]) == 5 assert clusters["1"]["id"] == "1" assert clusters["1"]["cluster_label"] == "my_nested_cluster" assert clusters["1"]["node_label"] == "my_other_info_node" - assert clusters["1"]["parent_id"] == "0" + assert clusters["1"]["cluster_id"] == "0" def test_get_clusters_doesnt_mutate(self): """Tests that get_clusters doesn't mutate state""" From a364869a27f8a78490be34403fe30038af4fd24f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:22:41 -0500 Subject: [PATCH 067/245] cleanup --- .../python_interface/visualization/construct_circuit_dag.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 419943c361..d497a95d2e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -17,10 +17,9 @@ from functools import singledispatchmethod from typing import Any -from xdsl.dialects import builtin, scf +from xdsl.dialects import builtin from xdsl.ir import Block, Operation, Region -from catalyst.python_interface.dialects import quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder From cddb234a7b3c20d7d6657e0e4aebfdfb5468b80a Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:24:28 -0500 Subject: [PATCH 068/245] clean-up --- .../visualization/construct_circuit_dag.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index d497a95d2e..0b6765ba3b 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -24,11 +24,11 @@ class ConstructCircuitDAG: - """A tool that analyzes an xDSL module and constructs a Directed Acyclic Graph (DAG) - using an injected DAGBuilder instance. This tool does not mutate the xDSL module.""" + """A tool that traverses an xDSL module and constructs a Directed Acyclic Graph (DAG) + of it's quantum program using an injected DAGBuilder instance. This tool does not mutate the xDSL module.""" def __init__(self, dag_builder: DAGBuilder) -> None: - """Initialize the analysis pass by injecting the DAG builder dependency. + """Initialize the utility by injecting the DAG builder dependency. Args: dag_builder (DAGBuilder): The concrete builder instance used for graph construction. @@ -45,7 +45,12 @@ def _visit(self, op: Any) -> None: to the specialized handler registered for its type.""" def construct(self, module: builtin.ModuleOp) -> None: - """Constructs the DAG from the module.""" + """Constructs the DAG from the module. + + Args: + module (xdsl.builtin.ModuleOp): The module containing the quantum program to visualize. + + """ for op in module.ops: self._visit(op) From 1b5210c512192826b1c8e5e61162eda0fc108bd4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:26:05 -0500 Subject: [PATCH 069/245] fix formatting issue --- .../python_interface/visualization/construct_circuit_dag.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0b6765ba3b..45a15cc2d2 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -25,7 +25,8 @@ class ConstructCircuitDAG: """A tool that traverses an xDSL module and constructs a Directed Acyclic Graph (DAG) - of it's quantum program using an injected DAGBuilder instance. This tool does not mutate the xDSL module.""" + of it's quantum program using an injected DAGBuilder instance. This tool does not mutate the xDSL module. + """ def __init__(self, dag_builder: DAGBuilder) -> None: """Initialize the utility by injecting the DAG builder dependency. From 14b28bb5fb2be009c3ea237dce271e3cc4942de5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:28:48 -0500 Subject: [PATCH 070/245] isort --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 5a55cd0ba7..ae449bb4c5 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -19,15 +19,16 @@ pytestmark = pytest.mark.usefixtures("requires_xdsl") +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class TestInitialization: From a1e9211922769cbd3ca1bd4276d945a81315cd98 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:40:31 -0500 Subject: [PATCH 071/245] update tests --- .../test_construct_circuit_dag.py | 97 +++++++++++++++---- 1 file changed, 78 insertions(+), 19 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index ae449bb4c5..a4b9f63765 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -19,46 +19,105 @@ pytestmark = pytest.mark.usefixtures("requires_xdsl") -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region -class TestInitialization: - """Tests that the state is correctly initialized.""" - - def test_dependency_injection(self): - """Tests that relevant dependencies are injected.""" - - mock_dag_builder = Mock(DAGBuilder) - utility = ConstructCircuitDAG(mock_dag_builder) - assert utility.dag_builder is mock_dag_builder - - +class FakeDAGBuilder(DAGBuilder): + """ + A concrete implementation of DAGBuilder used ONLY for testing. + It stores all graph manipulation calls in simple Python dictionaries + for easy assertion of the final graph state. + """ + + def __init__(self): + self._nodes = {} + self._edges = [] + self._clusters = {} + + def add_node(self, id, label, cluster_id=None, **attrs) -> None: + cluster_id = "__base__" if cluster_id is None else cluster_id + self._nodes[id] = { + "id": id, + "label": label, + "cluster_id": cluster_id, + "attrs": attrs, + } + + def add_edge(self, from_id: str, to_id: str, **attrs) -> None: + self._edges.append( + { + "from": from_id, + "to": to_id, + "attrs": attrs, + } + ) + + def add_cluster( + self, + id, + node_label=None, + cluster_id=None, + **attrs, + ) -> None: + cluster_id = "__base__" if cluster_id is None else cluster_id + self._clusters[id] = { + "id": id, + "label": node_label, + "cluster_id": cluster_id, + "attrs": attrs, + } + + def get_nodes(self): + return self._nodes.copy() + + def get_edges(self): + return self._edges.copy() + + def get_clusters(self): + return self._clusters.copy() + + def to_file(self, output_filename): + pass + + def to_string(self) -> str: + return "graph" + + +@pytest.mark.unit +def test_dependency_injection(): + """Tests that relevant dependencies are injected.""" + + dag_builder = FakeDAGBuilder() + utility = ConstructCircuitDAG(dag_builder) + assert utility.dag_builder is dag_builder + + +@pytest.mark.unit def test_does_not_mutate_module(): """Test that the module is not mutated.""" - # Create block containing some ops + # Create module op = test.TestOp() block = Block(ops=[op]) - # Create region containing some blocks region = Region(blocks=[block]) - # Create op containing the regions container_op = test.TestOp(regions=[region]) - # Create module op to house it all module_op = ModuleOp(ops=[container_op]) + # Save state before module_op_str_before = str(module_op) + # Process module mock_dag_builder = Mock(DAGBuilder) utility = ConstructCircuitDAG(mock_dag_builder) utility.construct(module_op) + # Ensure not mutated assert str(module_op) == module_op_str_before From aad7449103d4baf71a5213a92f21b34fd01832bb Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 25 Nov 2025 09:42:11 -0500 Subject: [PATCH 072/245] Apply suggestion from @andrijapau --- .../visualization/test_construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a4b9f63765..d9d677c153 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -33,7 +33,7 @@ class FakeDAGBuilder(DAGBuilder): """ A concrete implementation of DAGBuilder used ONLY for testing. - It stores all graph manipulation calls in simple Python dictionaries + It stores all graph manipulation calls in data structures for easy assertion of the final graph state. """ From 8225658dde58941bf298985fef93798dc3448315 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 09:45:07 -0500 Subject: [PATCH 073/245] isort --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index d9d677c153..cd44eef0cb 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -19,15 +19,16 @@ pytestmark = pytest.mark.usefixtures("requires_xdsl") +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 897fd1670cb47da18d87557d27dfe7438525655d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 13:56:07 -0500 Subject: [PATCH 074/245] feat: add bounding box visualization for funcops and devices --- .../visualization/construct_circuit_dag.py | 45 ++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 45a15cc2d2..c73a063e69 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -17,9 +17,10 @@ from functools import singledispatchmethod from typing import Any -from xdsl.dialects import builtin +from xdsl.dialects import builtin, func from xdsl.ir import Block, Operation, Region +from catalyst.python_interface.dialects import quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder @@ -36,6 +37,14 @@ def __init__(self, dag_builder: DAGBuilder) -> None: """ self.dag_builder: DAGBuilder = dag_builder + # Keep track of nesting clusters using a stack + # NOTE: `None` corresponds to the base graph 'cluster' + self._cluster_stack: list[str | None] = [None] + + def _reset(self) -> None: + """Resets the instance.""" + self._cluster_stack: list[str | None] = [None] + # ================================= # 1. CORE DISPATCH AND ENTRY POINT # ================================= @@ -62,6 +71,18 @@ def construct(self, module: builtin.ModuleOp) -> None: @_visit.register def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" + + # Visualize FuncOp's as bounding boxes + if isinstance(operation, func.FuncOp): + print(operation) + cluster_id = f"cluster_{id(operation)}" + self.dag_builder.add_cluster( + cluster_id, + node_label=operation.sym_name.data, + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) + for region in operation.regions: self._visit(region) @@ -76,3 +97,25 @@ def _block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: self._visit(op) + + @_visit.register + def _device_init(self, op: quantum.DeviceInitOp) -> None: + """Handles the initialization of a quantum device.""" + node_id = f"node_{id(op)}" + self.dag_builder.add_node( + node_id, + label=op.device_name.data, + cluster_id=self._cluster_stack[-1], + fillcolor="white", + shape="rectangle", + ) + + @_visit.register + def _func_return(self, op: func.ReturnOp) -> None: + """Handle func.return to exit FuncOp's cluster scope.""" + + # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp + if len(self._cluster_stack) > 2: + # If we hit a func.return operation we know we are leaving + # the FuncOp's scope and so we can pop the ID off the stack. + self._cluster_stack.pop() From 8ae869283bf61abe651b1c2adc5019b53657b2bd Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 25 Nov 2025 14:17:59 -0500 Subject: [PATCH 075/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index c73a063e69..c2321eaea3 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -74,7 +74,6 @@ def _operation(self, operation: Operation) -> None: # Visualize FuncOp's as bounding boxes if isinstance(operation, func.FuncOp): - print(operation) cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, From 3750ef2d650f01dbbe03f5fdedabb60804b1f631 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:21:18 -0500 Subject: [PATCH 076/245] cl --- doc/releases/changelog-dev.md | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 06f1966f41..1f4944ac71 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -6,6 +6,7 @@ [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213) [(#2229)](https://github.com/PennyLaneAI/catalyst/pull/2229) [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214) + [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From f823edeb58ddcf36e38d00417fcd145c1fc2a159 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:29:11 -0500 Subject: [PATCH 077/245] add tests for bounding boxes --- .../test_construct_circuit_dag.py | 65 +++++++++++++++++-- 1 file changed, 61 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index cd44eef0cb..aeaeaac202 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,19 +16,21 @@ from unittest.mock import Mock import pytest +from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors +import pennylane as qml +from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -122,3 +124,58 @@ def test_does_not_mutate_module(): # Ensure not mutated assert str(module_op) == module_op_str_before + + +@pytest.mark.unit +class TestFuncOpVisualization: + """Tests the visualization of FuncOps with bounding boxes""" + + def test_standard_qnode(self): + """Tests that a standard QJIT'd QNode is visualized correctly""" + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + qml.H(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.get_clusters() + # 1 = "jit_my_workflow" + # 2 = "my_workflow" + assert len(clusters) == 2 + + def test_nested_qnodes(self): + """Tests that nested QJIT'd QNodes are visualized correctly""" + + dev = qml.device("null.qubit", wires=1) + + @qml.qnode(dev) + def my_qnode2(): + qml.X(0) + + @qml.qnode(dev) + def my_qnode1(): + qml.H(0) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + def my_workflow(): + my_qnode1() + my_qnode2() + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.get_clusters() + # 1 = "jit_my_workflow" + # 2 = "my_qnode1" + # 3 = "my_qnode2" + assert len(clusters) == 3 From f397b4750c4c64e70652f2b9c25adee3a012707c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:30:27 -0500 Subject: [PATCH 078/245] isort --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index aeaeaac202..2137f97e54 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,14 +23,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 59e7826123dd0240d7347bc66aae77392b1e7a62 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:56:26 -0500 Subject: [PATCH 079/245] fix source code --- .../python_interface/visualization/construct_circuit_dag.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index c2321eaea3..13cbed1ff2 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -73,7 +73,10 @@ def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" # Visualize FuncOp's as bounding boxes - if isinstance(operation, func.FuncOp): + if isinstance(operation, func.FuncOp) and operation.sym_name.data not in { + "setup", + "teardown", + }: cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, From 55fbb8e7651f1b43eb8fb408aa6e44ec02c1cf28 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:57:50 -0500 Subject: [PATCH 080/245] Revert "fix source code" This reverts commit 59e7826123dd0240d7347bc66aae77392b1e7a62. --- .../python_interface/visualization/construct_circuit_dag.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 13cbed1ff2..c2321eaea3 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -73,10 +73,7 @@ def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" # Visualize FuncOp's as bounding boxes - if isinstance(operation, func.FuncOp) and operation.sym_name.data not in { - "setup", - "teardown", - }: + if isinstance(operation, func.FuncOp): cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, From 0e1b68bda28b530a9ad5a8ba90457f56a9c63d28 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 14:59:01 -0500 Subject: [PATCH 081/245] fix test --- .../visualization/test_construct_circuit_dag.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 2137f97e54..e9e7cea1d6 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -149,7 +149,9 @@ def my_workflow(): clusters = utility.dag_builder.get_clusters() # 1 = "jit_my_workflow" # 2 = "my_workflow" - assert len(clusters) == 2 + # 3 = "setup" + # 4 = "teardown" + assert len(clusters) == 4 def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -179,4 +181,6 @@ def my_workflow(): # 1 = "jit_my_workflow" # 2 = "my_qnode1" # 3 = "my_qnode2" + # 4 = "setup" + # 5 = "teardown" assert len(clusters) == 3 From 1749a1e209a477cd551d840466d814a554a297ad Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 25 Nov 2025 15:00:24 -0500 Subject: [PATCH 082/245] Apply suggestion from @andrijapau --- .../visualization/test_construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index e9e7cea1d6..90d2090c4e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -183,4 +183,4 @@ def my_workflow(): # 3 = "my_qnode2" # 4 = "setup" # 5 = "teardown" - assert len(clusters) == 3 + assert len(clusters) == 5 From 0b65d422ecfbdf797f487ef479bdd1b39a453e84 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 15:00:58 -0500 Subject: [PATCH 083/245] remove incorrect import --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 90d2090c4e..3496190ffe 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,7 +16,6 @@ from unittest.mock import Mock import pytest -from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") From 8f579a0b178307417dbd7b712dc1e4df9ff7f9a5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 15:39:52 -0500 Subject: [PATCH 084/245] update tests --- .../test_construct_circuit_dag.py | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3496190ffe..c51a8cf3f5 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -145,12 +144,17 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.get_clusters() - # 1 = "jit_my_workflow" - # 2 = "my_workflow" - # 3 = "setup" - # 4 = "teardown" - assert len(clusters) == 4 + graph_clusters = utility.dag_builder.get_clusters() + expected_cluster_labels = [ + "jit_my_workflow", + "my_workflow", + "setup", + "teardown", + ] + assert len(graph_clusters) == len(expected_cluster_labels) + cluster_labels = {info["label"] for info in graph_clusters.values()} + for expected_name in expected_cluster_labels: + assert expected_name in cluster_labels def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -176,10 +180,15 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.get_clusters() - # 1 = "jit_my_workflow" - # 2 = "my_qnode1" - # 3 = "my_qnode2" - # 4 = "setup" - # 5 = "teardown" - assert len(clusters) == 5 + graph_clusters = utility.dag_builder.get_clusters() + expected_cluster_labels = [ + "jit_my_workflow", + "my_qnode1", + "my_qnode2", + "setup", + "teardown", + ] + assert len(graph_clusters) == len(expected_cluster_labels) + cluster_labels = {info["label"] for info in graph_clusters.values()} + for expected_name in expected_cluster_labels: + assert expected_name in cluster_labels From 41056978825b0b6a12f4088a06357d56917ea223 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 15:54:22 -0500 Subject: [PATCH 085/245] isort --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index c51a8cf3f5..1af0e3813e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 028098e43718718b550a8bbcddf16bc3e8d18b67 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:02:19 -0500 Subject: [PATCH 086/245] minor visual change --- .../python_interface/visualization/construct_circuit_dag.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index c2321eaea3..e730c949bb 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -105,7 +105,9 @@ def _device_init(self, op: quantum.DeviceInitOp) -> None: node_id, label=op.device_name.data, cluster_id=self._cluster_stack[-1], - fillcolor="white", + fillcolor="grey", + color="black", + penwidth=2, shape="rectangle", ) From 9f373792e103928fc06a0376aa332b244e7b3ad6 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:16:33 -0500 Subject: [PATCH 087/245] feat: visualize control flow as clusters in the graph --- .../visualization/construct_circuit_dag.py | 67 ++++++++++++++++++- .../test_construct_circuit_dag.py | 43 ++++++++++-- 2 files changed, 103 insertions(+), 7 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index e730c949bb..2e10c25924 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -17,7 +17,7 @@ from functools import singledispatchmethod from typing import Any -from xdsl.dialects import builtin, func +from xdsl.dialects import builtin, func, scf from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import quantum @@ -82,8 +82,32 @@ def _operation(self, operation: Operation) -> None: ) self._cluster_stack.append(cluster_id) - for region in operation.regions: - self._visit(region) + if isinstance(operation, scf.IfOp): + # Loop through each branch and visualize as a cluster + for i, branch in enumerate(operation.regions): + cluster_id = f"cluster_ifop_branch{i}_{id(operation)}" + self.dag_builder.add_cluster( + cluster_id, + label=f"if ..." if i == 0 else "else", + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) + + # Go recursively into the branch to process internals + self._visit(branch) + + # Pop branch cluster after processing to ensure + # logical branches are treated as 'parallel' + self._cluster_stack.pop() + else: + for region in operation.regions: + self._visit(region) + + # Pop if the operation was a cluster creating operation + # This ensures proper nesting + ControlFlowOp = scf.ForOp | scf.WhileOp | scf.IfOp + if isinstance(operation, ControlFlowOp): + self._cluster_stack.pop() @_visit.register def _region(self, region: Region) -> None: @@ -120,3 +144,40 @@ def _func_return(self, op: func.ReturnOp) -> None: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. self._cluster_stack.pop() + + # ========================= + # 3. CONTROL FLOW HANDLERS + # ========================= + + @_visit.register + def _for_op(self, op: scf.ForOp) -> None: + """Handle an xDSL ForOp operation.""" + cluster_id = f"cluster_{id(op)}" + self.dag_builder.add_cluster( + cluster_id, + node_label="for ...", + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) + + @_visit.register + def _while_op(self, op: scf.WhileOp) -> None: + """Handle an xDSL WhileOp operation.""" + cluster_id = f"cluster_{id(op)}" + self.dag_builder.add_cluster( + cluster_id, + node_label="while ...", + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) + + @_visit.register + def _if_op(self, op: scf.IfOp) -> None: + """Handle an xDSL IfOp operation.""" + cluster_id = f"cluster_{id(op)}" + self.dag_builder.add_cluster( + cluster_id, + node_label="if", + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 1af0e3813e..9fdc0ad752 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -193,3 +192,39 @@ def my_workflow(): cluster_labels = {info["label"] for info in graph_clusters.values()} for expected_name in expected_cluster_labels: assert expected_name in cluster_labels + + +class TestControlFlowClusterVisualization: + """Tests that the control flow operations are visualized correctly as clusters.""" + + @pytest.mark.unit + def test_for_loop(self): + """Test that the for loop is visualized correctly.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + for i in range(3): + qml.H(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.get_clusters() + cluster_labels = {info["label"] for info in clusters.values()} + assert "for" in cluster_labels + + @pytest.mark.unit + def test_while_loop(self): + """Test that the while loop is visualized correctly.""" + pass + + @pytest.mark.unit + def test_conditional(self): + """Test that the conditional operation is visualized correctly.""" + pass From a5b8cd24892dc6a4baf06082813f712834dc141b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:16:44 -0500 Subject: [PATCH 088/245] format --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 9fdc0ad752..10bf458281 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 07a059774729df3cd328c7642d1a222e8ecb1b71 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 25 Nov 2025 16:25:44 -0500 Subject: [PATCH 089/245] Update frontend/catalyst/python_interface/visualization/construct_circuit_dag.py Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- .../python_interface/visualization/construct_circuit_dag.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 45a15cc2d2..0526e70b0d 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -29,11 +29,6 @@ class ConstructCircuitDAG: """ def __init__(self, dag_builder: DAGBuilder) -> None: - """Initialize the utility by injecting the DAG builder dependency. - - Args: - dag_builder (DAGBuilder): The concrete builder instance used for graph construction. - """ self.dag_builder: DAGBuilder = dag_builder # ================================= From 40b1eb4885996f505dfdeeec1719c788946aa2f7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:26:30 -0500 Subject: [PATCH 090/245] move things around --- .../visualization/construct_circuit_dag.py | 14 +++++--------- 1 file changed, 5 insertions(+), 9 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0526e70b0d..401a5d761c 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -31,15 +31,6 @@ class ConstructCircuitDAG: def __init__(self, dag_builder: DAGBuilder) -> None: self.dag_builder: DAGBuilder = dag_builder - # ================================= - # 1. CORE DISPATCH AND ENTRY POINT - # ================================= - - @singledispatchmethod - def _visit(self, op: Any) -> None: - """Central dispatch method (Visitor Pattern). Routes the operation 'op' - to the specialized handler registered for its type.""" - def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. @@ -54,6 +45,11 @@ def construct(self, module: builtin.ModuleOp) -> None: # 2. IR TRAVERSAL # ======================= + @singledispatchmethod + def _visit(self, op: Any) -> None: + """Central dispatch method (Visitor Pattern). Routes the operation 'op' + to the specialized handler registered for its type.""" + @_visit.register def _operation(self, operation: Operation) -> None: """Visit an xDSL Operation.""" From 77dd502c7cb8765d269a265f42a2fcd0cd7a3899 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:27:36 -0500 Subject: [PATCH 091/245] minor cleanup --- .../python_interface/visualization/construct_circuit_dag.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 401a5d761c..92e0febaed 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -41,9 +41,9 @@ def construct(self, module: builtin.ModuleOp) -> None: for op in module.ops: self._visit(op) - # ======================= - # 2. IR TRAVERSAL - # ======================= + # ============= + # IR TRAVERSAL + # ============= @singledispatchmethod def _visit(self, op: Any) -> None: From 19af1c71958a981a9b50a2610428bcd54f5914ba Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 16:31:35 -0500 Subject: [PATCH 092/245] clean-up --- .../visualization/construct_circuit_dag.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 36bd5a6905..4a8679dea9 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -88,6 +88,10 @@ def _block(self, block: Block) -> None: for op in block.ops: self._visit(op) + # ============ + # DEVICE NODE + # ============ + @_visit.register def _device_init(self, op: quantum.DeviceInitOp) -> None: """Handles the initialization of a quantum device.""" @@ -102,6 +106,10 @@ def _device_init(self, op: quantum.DeviceInitOp) -> None: shape="rectangle", ) + # ======================= + # FuncOp NESTING UTILITY + # ======================= + @_visit.register def _func_return(self, op: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" From b4ccd611bf89a3b9efc3ddc1b7d2a3099542de0c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 17:08:32 -0500 Subject: [PATCH 093/245] refactor the get_ to properties --- .../visualization/dag_builder.py | 11 ++- .../visualization/pydot_dag_builder.py | 15 ++-- .../visualization/test_dag_builder.py | 15 ++-- .../visualization/test_pydot_dag_builder.py | 73 +++---------------- 4 files changed, 35 insertions(+), 79 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 8a2fff6960..a268e24351 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -13,7 +13,7 @@ # limitations under the License. """File that defines the DAGBuilder abstract base class.""" -from abc import ABC, abstractmethod +from abc import ABC, abstractmethod, abstractproperty from typing import Any, TypeAlias ClusterID: TypeAlias = str @@ -84,8 +84,9 @@ def add_cluster( """ raise NotImplementedError + @property @abstractmethod - def get_nodes(self) -> dict[NodeID, dict[str, Any]]: + def nodes(self) -> dict[NodeID, dict[str, Any]]: """Retrieve the current set of nodes in the graph. Returns: @@ -93,8 +94,9 @@ def get_nodes(self) -> dict[NodeID, dict[str, Any]]: """ raise NotImplementedError + @property @abstractmethod - def get_edges(self) -> list[dict[str, Any]]: + def edges(self) -> list[dict[str, Any]]: """Retrieve the current set of edges in the graph. Returns: @@ -102,8 +104,9 @@ def get_edges(self) -> list[dict[str, Any]]: """ raise NotImplementedError + @property @abstractmethod - def get_clusters(self) -> dict[ClusterID, dict[str, Any]]: + def clusters(self) -> dict[ClusterID, dict[str, Any]]: """Retrieve the current set of clusters in the graph. Returns: diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index e8abda83da..77ee6d9fd7 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -204,29 +204,32 @@ def add_cluster( "attrs": dict(cluster_attrs), } - def get_nodes(self) -> dict[str, dict[str, Any]]: + @property + def nodes(self) -> dict[str, dict[str, Any]]: """Retrieve the current set of nodes in the graph. Returns: nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. """ - return self._nodes.copy() + return self._nodes - def get_edges(self) -> list[dict[str, Any]]: + @property + def edges(self) -> list[dict[str, Any]]: """Retrieve the current set of edges in the graph. Returns: edges (list[dict[str, Any]]): A list of edges where each element in the list contains a dictionary of edge information. """ - return self._edges.copy() + return self._edges - def get_clusters(self) -> dict[str, dict[str, Any]]: + @property + def clusters(self) -> dict[str, dict[str, Any]]: """Retrieve the current set of clusters in the graph. Returns: clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. """ - return self._clusters.copy() + return self._clusters def to_file(self, output_filename: str) -> None: """Save the graph to a file. diff --git a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py index 2ff5b8255f..df3a431ae2 100644 --- a/frontend/test/pytest/python_interface/visualization/test_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_dag_builder.py @@ -52,13 +52,16 @@ def add_cluster( ) -> None: return - def get_nodes(self) -> dict[str, dict[str, Any]]: + @property + def nodes(self) -> dict[str, dict[str, Any]]: return {} - def get_edges(self) -> list[dict[str, Any]]: + @property + def edges(self) -> list[dict[str, Any]]: return [] - def get_clusters(self) -> dict[str, dict[str, Any]]: + @property + def clusters(self) -> dict[str, dict[str, Any]]: return {} def to_file(self, output_filename: str) -> None: @@ -72,9 +75,9 @@ def to_string(self) -> str: node = dag_builder.add_node("0", "node0") edge = dag_builder.add_edge("0", "1") cluster = dag_builder.add_cluster("0") - nodes = dag_builder.get_nodes() - edges = dag_builder.get_edges() - clusters = dag_builder.get_clusters() + nodes = dag_builder.nodes + edges = dag_builder.edges + clusters = dag_builder.clusters render = dag_builder.to_file("test.png") string = dag_builder.to_string() diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 761be2da21..4833c0cc2c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -243,18 +243,18 @@ def test_add_cluster_with_attrs(self): assert cluster2.get("fontname") == "Helvetica" -class TestGetMethods: - """Tests the get_* methods.""" +class TestProperties: + """Tests the properties.""" - def test_get_nodes(self): - """Tests that get_nodes works.""" + def test_nodes(self): + """Tests that nodes works.""" dag_builder = PyDotDAGBuilder() dag_builder.add_node("0", "node0", fillcolor="red") dag_builder.add_cluster("c0") dag_builder.add_node("1", "node1", cluster_id="c0") - nodes = dag_builder.get_nodes() + nodes = dag_builder.nodes assert len(nodes) == 2 assert len(nodes["0"]) == 4 @@ -268,32 +268,15 @@ def test_get_nodes(self): assert nodes["1"]["label"] == "node1" assert nodes["1"]["cluster_id"] == "c0" - def test_get_nodes_doesnt_mutate(self): - """Tests that get_nodes doesn't mutate state""" - - dag_builder = PyDotDAGBuilder() - - dag_builder.add_node("0", "node0") - - old_nodes = dag_builder.get_nodes() - - dag_builder.add_node("1", "node1") - - new_nodes = dag_builder.get_nodes() - - assert old_nodes is not new_nodes - assert len(old_nodes) == 1 - assert len(new_nodes) == 2 - - def test_get_edges(self): - """Tests that get_edges works.""" + def test_edges(self): + """Tests that edges works.""" dag_builder = PyDotDAGBuilder() dag_builder.add_node("0", "node0") dag_builder.add_node("1", "node1") dag_builder.add_edge("0", "1", penwidth=10) - edges = dag_builder.get_edges() + edges = dag_builder.edges assert len(edges) == 1 @@ -301,37 +284,18 @@ def test_get_edges(self): assert edges[0]["to_id"] == "1" assert edges[0]["attrs"]["penwidth"] == 10 - def test_get_edges_doesnt_mutate(self): - """Tests that get_edges doesn't mutated.""" - - dag_builder = PyDotDAGBuilder() - dag_builder.add_node("0", "node0") - dag_builder.add_node("1", "node1") - dag_builder.add_edge("0", "1") - - old_edges = dag_builder.get_edges() - - dag_builder.add_node("2", "node2") - dag_builder.add_edge("1", "2") - - new_edges = dag_builder.get_edges() - - assert old_edges is not new_edges - assert len(old_edges) == 1 - assert len(new_edges) == 2 - def test_get_clusters(self): """Tests that get_clusters works.""" dag_builder = PyDotDAGBuilder() dag_builder.add_cluster("0", "my_info_node", label="my_cluster", penwidth=10) - clusters = dag_builder.get_clusters() + clusters = dag_builder.clusters dag_builder.add_cluster( "1", "my_other_info_node", cluster_id="0", label="my_nested_cluster" ) - clusters = dag_builder.get_clusters() + clusters = dag_builder.clusters assert len(clusters) == 2 assert len(clusters["0"]) == 5 @@ -347,23 +311,6 @@ def test_get_clusters(self): assert clusters["1"]["node_label"] == "my_other_info_node" assert clusters["1"]["cluster_id"] == "0" - def test_get_clusters_doesnt_mutate(self): - """Tests that get_clusters doesn't mutate state""" - - dag_builder = PyDotDAGBuilder() - - dag_builder.add_cluster("0") - - old_clusters = dag_builder.get_clusters() - - dag_builder.add_cluster("1") - - new_clusters = dag_builder.get_clusters() - - assert old_clusters is not new_clusters - assert len(old_clusters) == 1 - assert len(new_clusters) == 2 - class TestOutput: """Test that the graph can be outputted correctly.""" From d14c15bd52da7cda2fb982d1b6f3dcc432e32b69 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 17:13:00 -0500 Subject: [PATCH 094/245] update fake dag builder --- .../visualization/test_construct_circuit_dag.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index cd44eef0cb..8e6bdaef7c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -76,14 +76,17 @@ def add_cluster( "attrs": attrs, } - def get_nodes(self): - return self._nodes.copy() + @property + def nodes(self): + return self._nodes - def get_edges(self): - return self._edges.copy() + @property + def edges(self): + return self._edges - def get_clusters(self): - return self._clusters.copy() + @property + def clusters(self): + return self._clusters def to_file(self, output_filename): pass From d4295669ed1235334ea6e6a2de29bc03b0faacee Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 17:13:30 -0500 Subject: [PATCH 095/245] use new properties --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index f9c6fb22c8..8159da5e5d 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -148,7 +148,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_clusters = utility.dag_builder.get_clusters() + graph_clusters = utility.dag_builder.clusters expected_cluster_labels = [ "jit_my_workflow", "my_workflow", @@ -184,7 +184,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_clusters = utility.dag_builder.get_clusters() + graph_clusters = utility.dag_builder.clusters expected_cluster_labels = [ "jit_my_workflow", "my_qnode1", From a1110ade2f61372f341f15196e4b8ee6f56b679f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 17:15:06 -0500 Subject: [PATCH 096/245] add test and use new properties --- .../test_construct_circuit_dag.py | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a25a7c31e3..2b792eb156 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -219,14 +218,32 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.get_clusters() + clusters = utility.dag_builder.clusters cluster_labels = {info["label"] for info in clusters.values()} - assert "for" in cluster_labels + assert "for ..." in cluster_labels @pytest.mark.unit def test_while_loop(self): """Test that the while loop is visualized correctly.""" - pass + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + counter = 0 + while counter < 5: + qml.H(0) + counter += 1 + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.clusters + cluster_labels = {info["label"] for info in clusters.values()} + assert "while..." in cluster_labels @pytest.mark.unit def test_conditional(self): From cbb96ff2b1ebb9b88be8f7dc89d21ab27c976ec4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 25 Nov 2025 17:16:35 -0500 Subject: [PATCH 097/245] isort --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 2b792eb156..dcded52545 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 9c0220b583417c2859f69b597c6e9e02d99aad8c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 09:15:25 -0500 Subject: [PATCH 098/245] fix test --- .../visualization/test_construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index dcded52545..bb7649ea80 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -244,7 +244,7 @@ def my_workflow(): clusters = utility.dag_builder.clusters cluster_labels = {info["label"] for info in clusters.values()} - assert "while..." in cluster_labels + assert "while ..." in cluster_labels @pytest.mark.unit def test_conditional(self): From 0e0ddfdc77044d2d35ed4aa0c21d6592221252c1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 09:24:58 -0500 Subject: [PATCH 099/245] attempt to get rid of _subgraphs --- .../visualization/pydot_dag_builder.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 77ee6d9fd7..9a2922f6ec 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -53,10 +53,6 @@ def __init__( self.graph: pydot.Dot = pydot.Dot( graph_type="digraph", rankdir="TB", compound="true", strict=True ) - # Create cache for easy look-up - # TODO: Get rid of this and use self._clusters if possible - self._subgraphs: dict[str, pydot.Graph] = {} - self._subgraphs["__base__"] = self.graph # Internal state for graph structure self._nodes: dict[str, dict[str, Any]] = {} @@ -115,9 +111,14 @@ def add_node( # Use ChainMap so you don't need to construct a new dictionary node_attrs: ChainMap = ChainMap(attrs, self._default_node_attrs) node = pydot.Node(id, label=label, **node_attrs) - cluster_id = "__base__" if cluster_id is None else cluster_id - self._subgraphs[cluster_id].add_node(node) + # Add node to cluster + if cluster_id is None: + self.graph.add_node(node) + else: + # Use cluster ID to look up the subgraph + assert len(self.graph.get_subgraph(cluster_id)) == 1 + self.graph.get_subgraph(cluster_id)[0].add_node(node) self._nodes[id] = { "id": id, @@ -191,10 +192,13 @@ def add_cluster( cluster.add_subgraph(rank_subgraph) cluster.add_node(node) - self._subgraphs[id] = cluster - - cluster_id = "__base__" if cluster_id is None else cluster_id - self._subgraphs[cluster_id].add_subgraph(cluster) + # Add cluster to parent cluster + if cluster_id is None: + self.graph.add_subgraph(cluster) + else: + # Use cluster ID to look up the subgraph + assert len(self.graph.get_subgraph(cluster_id)) == 1 + self.graph.get_subgraph(cluster_id)[0].add_subgraph(cluster) self._clusters[id] = { "id": id, From c41adb9c16b2cb8733ba6aa3ceb4879412605e66 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 09:25:47 -0500 Subject: [PATCH 100/245] clean-up test --- .../python_interface/visualization/test_pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 4833c0cc2c..36f25ea990 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -284,8 +284,8 @@ def test_edges(self): assert edges[0]["to_id"] == "1" assert edges[0]["attrs"]["penwidth"] == 10 - def test_get_clusters(self): - """Tests that get_clusters works.""" + def test_clusters(self): + """Tests that clusters property works.""" dag_builder = PyDotDAGBuilder() dag_builder.add_cluster("0", "my_info_node", label="my_cluster", penwidth=10) From e8035430b87243871fc40365b1687d9003115bb6 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 09:29:32 -0500 Subject: [PATCH 101/245] rename __base__ to None --- .../python_interface/visualization/test_pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 36f25ea990..436ebc8319 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -261,7 +261,7 @@ def test_nodes(self): assert nodes["0"]["id"] == "0" assert nodes["0"]["label"] == "node0" - assert nodes["0"]["cluster_id"] == "__base__" + assert nodes["0"]["cluster_id"] == None assert nodes["0"]["attrs"]["fillcolor"] == "red" assert nodes["1"]["id"] == "1" @@ -302,7 +302,7 @@ def test_clusters(self): assert clusters["0"]["id"] == "0" assert clusters["0"]["cluster_label"] == "my_cluster" assert clusters["0"]["node_label"] == "my_info_node" - assert clusters["0"]["cluster_id"] == "__base__" + assert clusters["0"]["cluster_id"] == None assert clusters["0"]["attrs"]["penwidth"] == 10 assert len(clusters["1"]) == 5 From 085ec57f34919d38792e8f3f362663576bd6727e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 09:59:35 -0500 Subject: [PATCH 102/245] clean-up --- .../visualization/pydot_dag_builder.py | 24 +++++++++++-------- 1 file changed, 14 insertions(+), 10 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 9a2922f6ec..f034131a91 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -22,6 +22,7 @@ has_pydot = True try: import pydot + from pydot import Cluster, Dot, Edge, Node, Subgraph except ImportError: has_pydot = False @@ -50,7 +51,7 @@ def __init__( # - rankdir="TB": Set layout direction from Top to Bottom. # - compound="true": Allow edges to connect directly to clusters/subgraphs. # - strict=True: Prevent duplicate edges (e.g., A -> B added twice). - self.graph: pydot.Dot = pydot.Dot( + self.graph: Dot = Dot( graph_type="digraph", rankdir="TB", compound="true", strict=True ) @@ -110,15 +111,16 @@ def add_node( """ # Use ChainMap so you don't need to construct a new dictionary node_attrs: ChainMap = ChainMap(attrs, self._default_node_attrs) - node = pydot.Node(id, label=label, **node_attrs) + node = Node(id, label=label, **node_attrs) # Add node to cluster if cluster_id is None: self.graph.add_node(node) else: # Use cluster ID to look up the subgraph - assert len(self.graph.get_subgraph(cluster_id)) == 1 - self.graph.get_subgraph(cluster_id)[0].add_node(node) + parent_clusters = self.graph.get_subgraph(cluster_id)[0] + assert len(parent_clusters) == 1 + parent_clusters[0].add_node(node) self._nodes[id] = { "id": id, @@ -138,7 +140,8 @@ def add_edge(self, from_id: str, to_id: str, **attrs: Any) -> None: """ # Use ChainMap so you don't need to construct a new dictionary edge_attrs: ChainMap = ChainMap(attrs, self._default_edge_attrs) - edge = pydot.Edge(from_id, to_id, **edge_attrs) + edge = Edge(from_id, to_id, **edge_attrs) + self.graph.add_edge(edge) self._edges.append( @@ -166,7 +169,7 @@ def add_cluster( """ # Use ChainMap so you don't need to construct a new dictionary cluster_attrs: ChainMap = ChainMap(attrs, self._default_cluster_attrs) - cluster = pydot.Cluster(graph_name=id, **cluster_attrs) + cluster = Cluster(graph_name=id, **cluster_attrs) # Puts the label in a node within the cluster. # Ensures that any edges connecting nodes through the cluster @@ -179,8 +182,8 @@ def add_cluster( # └───────────┘ if node_label: node_id = f"{cluster_id}_info_node" - rank_subgraph = pydot.Subgraph() - node = pydot.Node( + rank_subgraph = Subgraph() + node = Node( node_id, label=node_label, shape="rectangle", @@ -197,8 +200,9 @@ def add_cluster( self.graph.add_subgraph(cluster) else: # Use cluster ID to look up the subgraph - assert len(self.graph.get_subgraph(cluster_id)) == 1 - self.graph.get_subgraph(cluster_id)[0].add_subgraph(cluster) + parent_clusters = self.graph.get_subgraph(cluster_id)[0] + assert len(parent_clusters) == 1 + parent_clusters[0].add_subgraph(cluster) self._clusters[id] = { "id": id, From c08a84cc82b7ab58a215b977547e3922f0b134b2 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:01:27 -0500 Subject: [PATCH 103/245] whoops --- .../python_interface/visualization/pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index f034131a91..aad95a2faf 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -118,7 +118,7 @@ def add_node( self.graph.add_node(node) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph(cluster_id)[0] + parent_clusters = self.graph.get_subgraph(cluster_id) assert len(parent_clusters) == 1 parent_clusters[0].add_node(node) @@ -200,7 +200,7 @@ def add_cluster( self.graph.add_subgraph(cluster) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph(cluster_id)[0] + parent_clusters = self.graph.get_subgraph(cluster_id) assert len(parent_clusters) == 1 parent_clusters[0].add_subgraph(cluster) From edc20765161feb4af228518072b7f32f6f7aa0b7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:03:37 -0500 Subject: [PATCH 104/245] add cluster_ prefix --- .../python_interface/visualization/pydot_dag_builder.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index aad95a2faf..9e617460d0 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -118,7 +118,7 @@ def add_node( self.graph.add_node(node) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph(cluster_id) + parent_clusters = self.graph.get_subgraph("cluster_"+cluster_id) assert len(parent_clusters) == 1 parent_clusters[0].add_node(node) @@ -200,7 +200,7 @@ def add_cluster( self.graph.add_subgraph(cluster) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph(cluster_id) + parent_clusters = self.graph.get_subgraph("cluster_"+cluster_id) assert len(parent_clusters) == 1 parent_clusters[0].add_subgraph(cluster) From 4c71876b8e2f0bbe5eddb361b14a1756c36d0af1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:22:12 -0500 Subject: [PATCH 105/245] add debug string --- .../visualization/pydot_dag_builder.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 9e617460d0..4960a1507a 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -118,8 +118,10 @@ def add_node( self.graph.add_node(node) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph("cluster_"+cluster_id) - assert len(parent_clusters) == 1 + parent_clusters = self.graph.get_subgraph("cluster_" + cluster_id) + assert len(parent_clusters) == 1, ( + f"Found {len(parent_clusters)} parent clusters with id {'cluster_' + cluster_id}" + ) parent_clusters[0].add_node(node) self._nodes[id] = { @@ -169,7 +171,7 @@ def add_cluster( """ # Use ChainMap so you don't need to construct a new dictionary cluster_attrs: ChainMap = ChainMap(attrs, self._default_cluster_attrs) - cluster = Cluster(graph_name=id, **cluster_attrs) + cluster = Cluster(id, **cluster_attrs) # Puts the label in a node within the cluster. # Ensures that any edges connecting nodes through the cluster @@ -200,8 +202,10 @@ def add_cluster( self.graph.add_subgraph(cluster) else: # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph("cluster_"+cluster_id) - assert len(parent_clusters) == 1 + parent_clusters = self.graph.get_subgraph("cluster_" + cluster_id) + assert len(parent_clusters) == 1, ( + f"Found {len(parent_clusters)} parent clusters with id {'cluster_' + cluster_id}" + ) parent_clusters[0].add_subgraph(cluster) self._clusters[id] = { From f1e5849b202f7af05a99ad238a614952d532e54b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:32:42 -0500 Subject: [PATCH 106/245] bring back cache --- .../visualization/pydot_dag_builder.py | 21 +++++++------------ .../visualization/test_pydot_dag_builder.py | 8 +++---- 2 files changed, 11 insertions(+), 18 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 4960a1507a..740b00f378 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -22,7 +22,7 @@ has_pydot = True try: import pydot - from pydot import Cluster, Dot, Edge, Node, Subgraph + from pydot import Cluster, Dot, Edge, Graph, Node, Subgraph except ImportError: has_pydot = False @@ -55,6 +55,9 @@ def __init__( graph_type="digraph", rankdir="TB", compound="true", strict=True ) + # Use internal cache that maps cluster ID to actual pydot (Dot or Cluster) object + self._subgraph_cache: dict[str, Graph] = {} + # Internal state for graph structure self._nodes: dict[str, dict[str, Any]] = {} self._edges: list[dict[str, Any]] = [] @@ -117,12 +120,7 @@ def add_node( if cluster_id is None: self.graph.add_node(node) else: - # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph("cluster_" + cluster_id) - assert len(parent_clusters) == 1, ( - f"Found {len(parent_clusters)} parent clusters with id {'cluster_' + cluster_id}" - ) - parent_clusters[0].add_node(node) + parent_cluster = self._subgraph_cache[cluster_id].add_node(node) self._nodes[id] = { "id": id, @@ -197,16 +195,11 @@ def add_cluster( cluster.add_subgraph(rank_subgraph) cluster.add_node(node) - # Add cluster to parent cluster + # Add node to cluster if cluster_id is None: self.graph.add_subgraph(cluster) else: - # Use cluster ID to look up the subgraph - parent_clusters = self.graph.get_subgraph("cluster_" + cluster_id) - assert len(parent_clusters) == 1, ( - f"Found {len(parent_clusters)} parent clusters with id {'cluster_' + cluster_id}" - ) - parent_clusters[0].add_subgraph(cluster) + parent_cluster = self._subgraph_cache[cluster_id].add_node(cluster) self._clusters[id] = { "id": id, diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 436ebc8319..6474b42d68 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -119,13 +119,13 @@ def test_add_cluster_to_parent_graph(self): # Level 0 (Root): Adds cluster on top of base graph dag_builder.add_node("n_root", "node_root") - dag_builder.add_cluster("c0") - # Level 1 (Inside c0): Add node on outer cluster and create new cluster on top + # Level 1 (c0): Add node on outer cluster + dag_builder.add_cluster("c0") dag_builder.add_node("n_outer", "node_outer", cluster_id="c0") - dag_builder.add_cluster("c1", cluster_id="c0") - # Level 2 (Inside c1): Add node on second cluster + # Level 2 (c1): Add node on inner cluster + dag_builder.add_cluster("c1", cluster_id="c0") dag_builder.add_node("n_inner", "node_inner", cluster_id="c1") root_graph = dag_builder.graph From 3c9ca3ae18112853ad21a0a83312205a501c7073 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:34:27 -0500 Subject: [PATCH 107/245] add good dev comment --- .../python_interface/visualization/pydot_dag_builder.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 740b00f378..d620b94987 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -56,6 +56,8 @@ def __init__( ) # Use internal cache that maps cluster ID to actual pydot (Dot or Cluster) object + # NOTE: This is needed so we don't need to traverse the graph to find the relevant + # cluster object to modify self._subgraph_cache: dict[str, Graph] = {} # Internal state for graph structure From bc54dbf2fd54d55f3c1437527785f020e414fc9b Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 10:44:03 -0500 Subject: [PATCH 108/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/pydot_dag_builder.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index d620b94987..008bdd3f7b 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -197,6 +197,9 @@ def add_cluster( cluster.add_subgraph(rank_subgraph) cluster.add_node(node) + # Record new cluster + self._subgraph_cache[id] = cluster + # Add node to cluster if cluster_id is None: self.graph.add_subgraph(cluster) From ee57f404802930981d675bdbae0b3aa82c96b52e Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 10:44:17 -0500 Subject: [PATCH 109/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 008bdd3f7b..ed690eb2ce 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -204,7 +204,7 @@ def add_cluster( if cluster_id is None: self.graph.add_subgraph(cluster) else: - parent_cluster = self._subgraph_cache[cluster_id].add_node(cluster) + parent_cluster = self._subgraph_cache[cluster_id].add_subgraph(cluster) self._clusters[id] = { "id": id, From c34185874df2365396001ceefd024036ced0ab16 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 10:45:07 -0500 Subject: [PATCH 110/245] whoops --- .../python_interface/visualization/pydot_dag_builder.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index ed690eb2ce..ff15f32f43 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -122,7 +122,7 @@ def add_node( if cluster_id is None: self.graph.add_node(node) else: - parent_cluster = self._subgraph_cache[cluster_id].add_node(node) + self._subgraph_cache[cluster_id].add_node(node) self._nodes[id] = { "id": id, @@ -199,12 +199,12 @@ def add_cluster( # Record new cluster self._subgraph_cache[id] = cluster - + # Add node to cluster if cluster_id is None: self.graph.add_subgraph(cluster) else: - parent_cluster = self._subgraph_cache[cluster_id].add_subgraph(cluster) + self._subgraph_cache[cluster_id].add_subgraph(cluster) self._clusters[id] = { "id": id, From 0460cac9ad28ce6e0e17382b6cdf74ef00ab131c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:03:51 -0500 Subject: [PATCH 111/245] refactor singledispatch --- .../visualization/construct_circuit_dag.py | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 92e0febaed..0b0f960064 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -39,31 +39,24 @@ def construct(self, module: builtin.ModuleOp) -> None: """ for op in module.ops: - self._visit(op) + self._visit_operation(op) # ============= # IR TRAVERSAL # ============= @singledispatchmethod - def _visit(self, op: Any) -> None: - """Central dispatch method (Visitor Pattern). Routes the operation 'op' - to the specialized handler registered for its type.""" - - @_visit.register - def _operation(self, operation: Operation) -> None: - """Visit an xDSL Operation.""" + def _visit_operation(self, operation: Operation) -> None: + """Visit an xDSL Operation. Default to visiting each region contained in the operation.""" for region in operation.regions: - self._visit(region) + self._visit_region(region) - @_visit.register - def _region(self, region: Region) -> None: + def _visit_region(self, region: Region) -> None: """Visit an xDSL Region operation.""" for block in region.blocks: - self._visit(block) + self._visit_block(block) - @_visit.register - def _block(self, block: Block) -> None: + def _visit_block(self, block: Block) -> None: """Visit an xDSL Block operation, dispatching handling for each contained Operation.""" for op in block.ops: - self._visit(op) + self._visit_operation(op) From 1a6aa98f3fd0e3e1cfa91a17ac34de8ca152907c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:08:29 -0500 Subject: [PATCH 112/245] minor improvements --- .../visualization/construct_circuit_dag.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index b50cc0b2e1..6a941da60b 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -53,7 +53,7 @@ def construct(self, module: builtin.ModuleOp) -> None: # ============= # IR TRAVERSAL # ============= - + @singledispatchmethod def _visit_operation(self, operation: Operation) -> None: """Visit an xDSL Operation. Default to visiting each region contained in the operation.""" @@ -75,12 +75,12 @@ def _visit_block(self, block: Block) -> None: # ============ @_visit_operation.register - def _device_init(self, op: quantum.DeviceInitOp) -> None: + def _device_init(self, operation: quantum.DeviceInitOp) -> None: """Handles the initialization of a quantum device.""" - node_id = f"node_{id(op)}" + node_id = f"node_{id(operation)}" self.dag_builder.add_node( node_id, - label=op.device_name.data, + label=operation.device_name.data, cluster_id=self._cluster_stack[-1], fillcolor="grey", color="black", @@ -88,6 +88,9 @@ def _device_init(self, op: quantum.DeviceInitOp) -> None: shape="rectangle", ) + for region in operation.regions: + self._visit_region(region) + # ======================= # FuncOp NESTING UTILITY # ======================= @@ -103,7 +106,10 @@ def _func_op(self, operation: func.FuncOp) -> None: cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) - + + for region in operation.regions: + self._visit_region(region) + @_visit_operation.register def _func_return(self, op: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" From 39fe0b9b4a876225360d16ec727ca06632b58cf8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:11:04 -0500 Subject: [PATCH 113/245] whoops --- .../python_interface/visualization/construct_circuit_dag.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 6a941da60b..c6041425e2 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -111,7 +111,7 @@ def _func_op(self, operation: func.FuncOp) -> None: self._visit_region(region) @_visit_operation.register - def _func_return(self, op: func.ReturnOp) -> None: + def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp @@ -119,3 +119,6 @@ def _func_return(self, op: func.ReturnOp) -> None: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. self._cluster_stack.pop() + + for region in operation.regions: + self._visit_region(region) From 194433f24b509d539cfae7f43ccb85187e620750 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:43:17 -0500 Subject: [PATCH 114/245] format --- .../visualization/construct_circuit_dag.py | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 7a269d985b..f87bad9197 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -94,7 +94,7 @@ def _operation(self, operation: Operation) -> None: ControlFlowOp = scf.ForOp | scf.WhileOp | scf.IfOp if isinstance(operation, ControlFlowOp): self._cluster_stack.pop() - + @singledispatchmethod def _visit_operation(self, operation: Operation) -> None: """Visit an xDSL Operation. Default to visiting each region contained in the operation.""" @@ -114,7 +114,7 @@ def _visit_block(self, block: Block) -> None: # ============= # CONTROL FLOW # ============= - + @_visit_operation.register def _for_op(self, op: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" @@ -125,10 +125,10 @@ def _for_op(self, op: scf.ForOp) -> None: cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) - + for region in operation.regions: self._visit_region(region) - + self._cluster_stack.pop() @_visit_operation.register @@ -141,12 +141,12 @@ def _while_op(self, op: scf.WhileOp) -> None: cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) - + for region in operation.regions: self._visit_region(region) self._cluster_stack.pop() - + @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" @@ -157,7 +157,7 @@ def _if_op(self, operation: scf.IfOp): cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) - + # Loop through each branch and visualize as a cluster for i, branch in enumerate(operation.regions): cluster_id = f"cluster_ifop_branch{i}_{id(operation)}" @@ -174,9 +174,9 @@ def _if_op(self, operation: scf.IfOp): # Pop branch cluster after processing to ensure # logical branches are treated as 'parallel' self._cluster_stack.pop() - + self._cluster_stack.pop() - + # ============ # DEVICE NODE # ============ From 76025f3732f47e539e8d040e13eb4dda0870a4fc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:44:21 -0500 Subject: [PATCH 115/245] more clean-up --- .../visualization/construct_circuit_dag.py | 41 ------------------- 1 file changed, 41 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index f87bad9197..8721b8415f 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -54,47 +54,6 @@ def construct(self, module: builtin.ModuleOp) -> None: # IR TRAVERSAL # ============= - @_visit.register - def _operation(self, operation: Operation) -> None: - """Visit an xDSL Operation.""" - - # Visualize FuncOp's as bounding boxes - if isinstance(operation, func.FuncOp): - cluster_id = f"cluster_{id(operation)}" - self.dag_builder.add_cluster( - cluster_id, - node_label=operation.sym_name.data, - cluster_id=self._cluster_stack[-1], - ) - self._cluster_stack.append(cluster_id) - - if isinstance(operation, scf.IfOp): - # Loop through each branch and visualize as a cluster - for i, branch in enumerate(operation.regions): - cluster_id = f"cluster_ifop_branch{i}_{id(operation)}" - self.dag_builder.add_cluster( - cluster_id, - label=f"if ..." if i == 0 else "else", - cluster_id=self._cluster_stack[-1], - ) - self._cluster_stack.append(cluster_id) - - # Go recursively into the branch to process internals - self._visit(branch) - - # Pop branch cluster after processing to ensure - # logical branches are treated as 'parallel' - self._cluster_stack.pop() - else: - for region in operation.regions: - self._visit(region) - - # Pop if the operation was a cluster creating operation - # This ensures proper nesting - ControlFlowOp = scf.ForOp | scf.WhileOp | scf.IfOp - if isinstance(operation, ControlFlowOp): - self._cluster_stack.pop() - @singledispatchmethod def _visit_operation(self, operation: Operation) -> None: """Visit an xDSL Operation. Default to visiting each region contained in the operation.""" From ccceb2b57fe1ff3d48ff3b15c681f78dfecc1153 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:45:20 -0500 Subject: [PATCH 116/245] clobbered --- .../python_interface/visualization/construct_circuit_dag.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 8721b8415f..0e671d00e6 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -185,3 +185,6 @@ def _func_return(self, operation: func.ReturnOp) -> None: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. self._cluster_stack.pop() + + for region in operation.regions: + self._visit_region(region) From 0cbe9afa9fcfde84518f3a8c8c2efedaddeb5821 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 11:47:01 -0500 Subject: [PATCH 117/245] Update frontend/catalyst/python_interface/visualization/dag_builder.py Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- frontend/catalyst/python_interface/visualization/dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index a268e24351..4cfe44854a 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -13,7 +13,7 @@ # limitations under the License. """File that defines the DAGBuilder abstract base class.""" -from abc import ABC, abstractmethod, abstractproperty +from abc import ABC, abstractmethod from typing import Any, TypeAlias ClusterID: TypeAlias = str From 813ec0e64373df74475cc6f2b5ecda52d1cc93a5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:48:23 -0500 Subject: [PATCH 118/245] remove unused --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0e671d00e6..2adfa87b7e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -15,7 +15,6 @@ """Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module.""" from functools import singledispatchmethod -from typing import Any from xdsl.dialects import builtin, func, scf from xdsl.ir import Block, Operation, Region From 3b1d25e44be58627168a574cec8894481e6aca65 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:48:37 -0500 Subject: [PATCH 119/245] remove unused --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index c6041425e2..9e33523a00 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -15,7 +15,6 @@ """Contains the ConstructCircuitDAG tool for constructing a DAG from an xDSL module.""" from functools import singledispatchmethod -from typing import Any from xdsl.dialects import builtin, func from xdsl.ir import Block, Operation, Region From 96d3b3795243ea27b821506f38c9b8ce10696c92 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 11:50:56 -0500 Subject: [PATCH 120/245] fix pylint --- .../visualization/construct_circuit_dag.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 2adfa87b7e..d456750078 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -74,9 +74,9 @@ def _visit_block(self, block: Block) -> None: # ============= @_visit_operation.register - def _for_op(self, op: scf.ForOp) -> None: + def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" - cluster_id = f"cluster_{id(op)}" + cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, node_label="for ...", @@ -90,9 +90,9 @@ def _for_op(self, op: scf.ForOp) -> None: self._cluster_stack.pop() @_visit_operation.register - def _while_op(self, op: scf.WhileOp) -> None: + def _while_op(self, operation: scf.WhileOp) -> None: """Handle an xDSL WhileOp operation.""" - cluster_id = f"cluster_{id(op)}" + cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, node_label="while ...", @@ -108,7 +108,7 @@ def _while_op(self, op: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - cluster_id = f"cluster_{id(op)}" + cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, node_label="if", From 2e2a97c891bb5ac69ce05b0aae2bbd93d590ea11 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 12:06:52 -0500 Subject: [PATCH 121/245] more tests --- .../test_construct_circuit_dag.py | 53 ++++++++++++++++++- 1 file changed, 51 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index bb7649ea80..40e860c018 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -247,6 +247,55 @@ def my_workflow(): assert "while ..." in cluster_labels @pytest.mark.unit - def test_conditional(self): + def test_if_else_conditional(self): """Test that the conditional operation is visualized correctly.""" - pass + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + flag = 1 + if flag == 1: + qml.X(0) + else: + qml.Y(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.clusters + cluster_labels = {info["label"] for info in clusters.values()} + assert "if ..." in cluster_labels + assert "else" in cluster_labels + + @pytest.mark.unit + def test_if_elif_else_conditional(self): + """Test that the conditional operation is visualized correctly.""" + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + flag = 1 + if flag == 1: + qml.X(0) + elif flag == 2: + qml.Y(0) + else: + qml.Z(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.clusters + cluster_labels = [info["label"] for info in clusters.values()] + assert "if ..." in cluster_labels + assert cluster_labels.count("if ...") == 2 + assert "else" in cluster_labels + assert cluster_labels.count("else") == 2 From e7523295c075a55626f030192624bd0af6dc64f9 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 12:16:58 -0500 Subject: [PATCH 122/245] add more details to docstring --- .../visualization/construct_circuit_dag.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0b0f960064..6b7500562d 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -24,8 +24,18 @@ class ConstructCircuitDAG: - """A tool that traverses an xDSL module and constructs a Directed Acyclic Graph (DAG) + """Utility tool following the director pattern to build a DAG representation of a compiled quantum program. + + This tool traverses an xDSL module and constructs a Directed Acyclic Graph (DAG) of it's quantum program using an injected DAGBuilder instance. This tool does not mutate the xDSL module. + + **Example** + + >>> builder = PyDotDAGBuilder() + >>> director = ConstructCircuitDAG(builder) + >>> director.construct(module) + >>> director.dag_builder.to_string() + ... """ def __init__(self, dag_builder: DAGBuilder) -> None: From 123144f6dd3dc2f71ce33adf83ead58e44abd4f7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 13:27:24 -0500 Subject: [PATCH 123/245] clean up tests --- .../visualization/construct_circuit_dag.py | 2 +- .../test_construct_circuit_dag.py | 29 ++++++++++++++----- 2 files changed, 23 insertions(+), 8 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index efd4e409de..3d0fa712d2 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -111,7 +111,7 @@ def _func_op(self, operation: func.FuncOp) -> None: cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, - node_label=operation.sym_name.data, + label=operation.sym_name.data, cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 8159da5e5d..95d675e91a 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -57,8 +57,8 @@ def add_node(self, id, label, cluster_id=None, **attrs) -> None: def add_edge(self, from_id: str, to_id: str, **attrs) -> None: self._edges.append( { - "from": from_id, - "to": to_id, + "from_id": from_id, + "to_id": to_id, "attrs": attrs, } ) @@ -73,7 +73,8 @@ def add_cluster( cluster_id = "__base__" if cluster_id is None else cluster_id self._clusters[id] = { "id": id, - "label": node_label, + "cluster_label": attrs.get("label"), + "node_label": node_label, "cluster_id": cluster_id, "attrs": attrs, } @@ -149,16 +150,22 @@ def my_workflow(): utility.construct(module) graph_clusters = utility.dag_builder.clusters + + # Check labels we expected are there expected_cluster_labels = [ "jit_my_workflow", "my_workflow", "setup", "teardown", ] - assert len(graph_clusters) == len(expected_cluster_labels) - cluster_labels = {info["label"] for info in graph_clusters.values()} - for expected_name in expected_cluster_labels: - assert expected_name in cluster_labels + generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + for cluster_label in expected_cluster_labels: + assert cluster_label in generated_cluster_labels + + # Check nesting is correct + # graph + # └── jit_my_workflow + # └── my_workflow def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -185,6 +192,8 @@ def my_workflow(): utility.construct(module) graph_clusters = utility.dag_builder.clusters + + # Check labels we expected are there as clusters expected_cluster_labels = [ "jit_my_workflow", "my_qnode1", @@ -196,3 +205,9 @@ def my_workflow(): cluster_labels = {info["label"] for info in graph_clusters.values()} for expected_name in expected_cluster_labels: assert expected_name in cluster_labels + + # Check nesting is correct + # graph + # └── jit_my_workflow + # ├── my_qnode1 + # └── my_qnode2 From d0c7c9fc81239ac886d8bc24591056d064a7c3cc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 13:42:32 -0500 Subject: [PATCH 124/245] try to improve tests --- .../test_construct_circuit_dag.py | 47 ++++++++++++++++--- 1 file changed, 40 insertions(+), 7 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 95d675e91a..853ea63f8e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,21 +16,21 @@ from unittest.mock import Mock import pytest +from _typeshed import GenericPath pytestmark = pytest.mark.usefixtures("requires_xdsl") # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -46,7 +46,6 @@ def __init__(self): self._clusters = {} def add_node(self, id, label, cluster_id=None, **attrs) -> None: - cluster_id = "__base__" if cluster_id is None else cluster_id self._nodes[id] = { "id": id, "label": label, @@ -70,7 +69,6 @@ def add_cluster( cluster_id=None, **attrs, ) -> None: - cluster_id = "__base__" if cluster_id is None else cluster_id self._clusters[id] = { "id": id, "cluster_label": attrs.get("label"), @@ -158,7 +156,9 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + generated_cluster_labels = { + info["cluster_label"] for info in graph_clusters.values() + } for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -167,6 +167,18 @@ def my_workflow(): # └── jit_my_workflow # └── my_workflow + # Get the parent labels for each cluster and ensure they are what we expect. + parent_labels = ( + graph_clusters[child_cluster["cluster_id"]]["label"] + for child_cluster in graph_clusters.values() + if child_cluster.get("cluster_id") is not None + ) + cluster_label_to_parent_label: dict[str, str] = dict( + zip(tuple(generated_cluster_labels), parent_labels) + ) + assert cluster_label_to_parent_label["jit_my_workflow"] is None + assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" + def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -211,3 +223,24 @@ def my_workflow(): # └── jit_my_workflow # ├── my_qnode1 # └── my_qnode2 + + # Get the parent labels for each cluster and ensure they are what we expect. + parent_labels = ( + graph_clusters[child_cluster["cluster_id"]]["label"] + for child_cluster in graph_clusters.values() + if child_cluster.get("cluster_id") is not None + ) + cluster_label_to_parent_label: dict[str, str] = dict( + zip(tuple(cluster_labels), parent_labels) + ) + assert cluster_label_to_parent_label["jit_my_workflow"] is None + assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" + assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" + + +class TestDeviceNode: + """Tests that the device node is correctly visualized.""" + + def test_standard_qnode(self): + """Tests that a standard setup works.""" + pass From cd9d5f49df8f11a8a43b1b53a230898a0587d886 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 13:46:41 -0500 Subject: [PATCH 125/245] add device test --- .../test_construct_circuit_dag.py | 33 ++++++++++++++++++- 1 file changed, 32 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 853ea63f8e..dd889f7b67 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -15,6 +15,7 @@ from unittest.mock import Mock +from jax import util import pytest from _typeshed import GenericPath @@ -243,4 +244,34 @@ class TestDeviceNode: def test_standard_qnode(self): """Tests that a standard setup works.""" - pass + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + qml.H(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + graph_nodes = utility.dag_builder.nodes + + # Basic check for node + node_labels = {info["label"] for info in graph_nodes.values()} + assert "NullQubit" in node_labels + + # Ensure nesting is correct + graph_clusters = utility.dag_builder.clusters + parent_labels = ( + graph_clusters[child_cluster["cluster_id"]]["label"] + for child_cluster in graph_clusters.values() + if child_cluster.get("cluster_id") is not None + ) + cluster_label_to_parent_label: dict[str, str] = dict( + zip(tuple(node_labels), parent_labels) + ) + assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" From 6cda3bff6b4231eecadc56bf831c5c64c33ba75b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 13:46:51 -0500 Subject: [PATCH 126/245] format --- .../test_construct_circuit_dag.py | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index dd889f7b67..6d632031e6 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -15,23 +15,24 @@ from unittest.mock import Mock -from jax import util import pytest from _typeshed import GenericPath +from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -157,9 +158,7 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = { - info["cluster_label"] for info in graph_clusters.values() - } + generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -244,7 +243,7 @@ class TestDeviceNode: def test_standard_qnode(self): """Tests that a standard setup works.""" - + dev = qml.device("null.qubit", wires=1) @xdsl_from_qjit @@ -271,7 +270,5 @@ def my_workflow(): for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) - cluster_label_to_parent_label: dict[str, str] = dict( - zip(tuple(node_labels), parent_labels) - ) - assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" + cluster_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) + assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" From 86b566287f755c89305d298dddec10e326c738de Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:47:07 -0500 Subject: [PATCH 127/245] Update frontend/catalyst/python_interface/visualization/dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- frontend/catalyst/python_interface/visualization/dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index 4cfe44854a..ea1cb5438d 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -90,7 +90,7 @@ def nodes(self) -> dict[NodeID, dict[str, Any]]: """Retrieve the current set of nodes in the graph. Returns: - nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. + nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to its node information. """ raise NotImplementedError From 9c26efdc019ba3b4c258510827eebc8157d0b298 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:47:14 -0500 Subject: [PATCH 128/245] Update frontend/catalyst/python_interface/visualization/pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index ff15f32f43..06e9e00fb5 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -219,7 +219,7 @@ def nodes(self) -> dict[str, dict[str, Any]]: """Retrieve the current set of nodes in the graph. Returns: - nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to it's node information. + nodes (dict[str, dict[str, Any]]): A dictionary that maps the node's ID to its node information. """ return self._nodes From 81a9aa6969c8c937fa17aa55009fb327e0e96cd7 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:47:28 -0500 Subject: [PATCH 129/245] Update frontend/catalyst/python_interface/visualization/pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 06e9e00fb5..7fd23660cb 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -183,7 +183,7 @@ def add_cluster( # │ │ # └───────────┘ if node_label: - node_id = f"{cluster_id}_info_node" + node_id = f"{id}_info_node" rank_subgraph = Subgraph() node = Node( node_id, From ab06276ac071751e92e236c56036ba5e3292fdf8 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:47:44 -0500 Subject: [PATCH 130/245] Update frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/test_pydot_dag_builder.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 6474b42d68..f84fa80773 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -17,7 +17,6 @@ import pytest -from frontend.catalyst.python_interface.visualization import dag_builder pydot = pytest.importorskip("pydot") pytestmark = pytest.mark.usefixtures("requires_xdsl") From 3e4102b1139ff222644b95421c9e954aa0241f11 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:48:08 -0500 Subject: [PATCH 131/245] Update frontend/catalyst/python_interface/visualization/pydot_dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- .../python_interface/visualization/pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 7fd23660cb..5fd1cf7d74 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -237,7 +237,7 @@ def clusters(self) -> dict[str, dict[str, Any]]: """Retrieve the current set of clusters in the graph. Returns: - clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. + clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to its cluster information. """ return self._clusters From 194f14a55d86cb713584a7ac722105cb2f12bd48 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:48:16 -0500 Subject: [PATCH 132/245] Update frontend/catalyst/python_interface/visualization/dag_builder.py Co-authored-by: Mehrdad Malek <39844030+mehrdad2m@users.noreply.github.com> --- frontend/catalyst/python_interface/visualization/dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/dag_builder.py b/frontend/catalyst/python_interface/visualization/dag_builder.py index ea1cb5438d..70c5806616 100644 --- a/frontend/catalyst/python_interface/visualization/dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/dag_builder.py @@ -110,7 +110,7 @@ def clusters(self) -> dict[ClusterID, dict[str, Any]]: """Retrieve the current set of clusters in the graph. Returns: - clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to it's cluster information. + clusters (dict[str, dict[str, Any]]): A dictionary that maps the cluster's ID to its cluster information. """ raise NotImplementedError From 199ab705a17c851256da306fea7b20933e00fd50 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 13:54:00 -0500 Subject: [PATCH 133/245] add reset --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 3d0fa712d2..074bd50f03 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -56,6 +56,7 @@ def construct(self, module: builtin.ModuleOp) -> None: module (xdsl.builtin.ModuleOp): The module containing the quantum program to visualize. """ + self._reset() for op in module.ops: self._visit_operation(op) From 24d3dd3a807603072ee2fd536e6b6be73e0add61 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 13:59:59 -0500 Subject: [PATCH 134/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3de387ec46..8651ab12ff 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -217,7 +217,6 @@ def my_workflow(): cluster_labels = {info["label"] for info in graph_clusters.values()} for expected_name in expected_cluster_labels: assert expected_name in cluster_labels - # Check nesting is correct # graph # └── jit_my_workflow From 685842c5360c91705d08896ab28a535fb2dba7a0 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:02:38 -0500 Subject: [PATCH 135/245] fix fakebackend --- .../visualization/test_construct_circuit_dag.py | 9 ++++----- 1 file changed, 4 insertions(+), 5 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 8e6bdaef7c..4a98739c5a 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -44,7 +44,6 @@ def __init__(self): self._clusters = {} def add_node(self, id, label, cluster_id=None, **attrs) -> None: - cluster_id = "__base__" if cluster_id is None else cluster_id self._nodes[id] = { "id": id, "label": label, @@ -55,8 +54,8 @@ def add_node(self, id, label, cluster_id=None, **attrs) -> None: def add_edge(self, from_id: str, to_id: str, **attrs) -> None: self._edges.append( { - "from": from_id, - "to": to_id, + "from_id": from_id, + "to_id": to_id, "attrs": attrs, } ) @@ -68,10 +67,10 @@ def add_cluster( cluster_id=None, **attrs, ) -> None: - cluster_id = "__base__" if cluster_id is None else cluster_id self._clusters[id] = { "id": id, - "label": node_label, + "node_label": node_label, + "cluster_label": attrs.get("label"), "cluster_id": cluster_id, "attrs": attrs, } From 8c64d8187698a2288a0d9797d9046041c7251b6e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:04:05 -0500 Subject: [PATCH 136/245] isort --- .../python_interface/visualization/test_pydot_dag_builder.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index f84fa80773..b87b283903 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -17,7 +17,6 @@ import pytest - pydot = pytest.importorskip("pydot") pytestmark = pytest.mark.usefixtures("requires_xdsl") # pylint: disable=wrong-import-position From 9847b3948d1012fdd1d0a25e20b8cd6c6bc1d062 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:22:40 -0500 Subject: [PATCH 137/245] try moving cluster stack to dag builder --- .../visualization/construct_circuit_dag.py | 13 ++------- .../visualization/pydot_dag_builder.py | 28 +++++++++++++++++-- .../test_construct_circuit_dag.py | 2 -- 3 files changed, 29 insertions(+), 14 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 074bd50f03..924132c3da 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -41,13 +41,9 @@ class ConstructCircuitDAG: def __init__(self, dag_builder: DAGBuilder) -> None: self.dag_builder: DAGBuilder = dag_builder - # Keep track of nesting clusters using a stack - # NOTE: `None` corresponds to the base graph 'cluster' - self._cluster_stack: list[str | None] = [None] - def _reset(self) -> None: """Resets the instance.""" - self._cluster_stack: list[str | None] = [None] + self.dag_builder.reset() def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. @@ -91,7 +87,6 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: self.dag_builder.add_node( node_id, label=operation.device_name.data, - cluster_id=self._cluster_stack[-1], fillcolor="grey", color="black", penwidth=2, @@ -113,9 +108,7 @@ def _func_op(self, operation: func.FuncOp) -> None: self.dag_builder.add_cluster( cluster_id, label=operation.sym_name.data, - cluster_id=self._cluster_stack[-1], ) - self._cluster_stack.append(cluster_id) for region in operation.regions: self._visit_region(region) @@ -125,10 +118,10 @@ def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp - if len(self._cluster_stack) > 2: + if len(self.dag_builder.cluster_id_stack) > 2: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. - self._cluster_stack.pop() + self.dag_builder.cluster_id_stack.pop() for region in operation.regions: self._visit_region(region) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 5fd1cf7d74..5191ab23ba 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -65,6 +65,9 @@ def __init__( self._edges: list[dict[str, Any]] = [] self._clusters: dict[str, dict[str, Any]] = {} + # Keep track of nesting clusters using an internal stack + self.cluster_id_stack: list[str] = [] + _default_attrs: dict = ( {"fontname": "Helvetica", "penwidth": 2} if attrs is None else attrs ) @@ -98,6 +101,20 @@ def __init__( else cluster_attrs ) + def reset(self) -> None: + """Resets the instance.""" + self.graph: Dot = Dot( + graph_type="digraph", rankdir="TB", compound="true", strict=True + ) + + self._subgraph_cache: dict[str, Graph] = {} + + self._nodes: dict[str, dict[str, Any]] = {} + self._edges: list[dict[str, Any]] = [] + self._clusters: dict[str, dict[str, Any]] = {} + + self.cluster_id_stack: list[str] = [] + def add_node( self, id: str, @@ -200,12 +217,19 @@ def add_cluster( # Record new cluster self._subgraph_cache[id] = cluster - # Add node to cluster if cluster_id is None: - self.graph.add_subgraph(cluster) + if self.cluster_id_stack: + # Nest on top of top most cluster + self._subgraph_cache[self.cluster_id_stack[-1]].add_subgraph(cluster) + else: + # If stack is empty, place on base graph + self.graph.add_subgraph(cluster) else: self._subgraph_cache[cluster_id].add_subgraph(cluster) + # Add cluster to stack + self.cluster_id_stack.append(id) + self._clusters[id] = { "id": id, "cluster_label": cluster_attrs.get("label"), diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a210a104ef..55e6116324 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,8 +16,6 @@ from unittest.mock import Mock import pytest -from _typeshed import GenericPath -from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") From 57ce573ea2369e60c2c5eff576b6c28bb3f8e370 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:24:44 -0500 Subject: [PATCH 138/245] add test --- .../visualization/test_construct_circuit_dag.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 55e6116324..ce8b9a1b02 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -128,6 +128,12 @@ def test_does_not_mutate_module(): assert str(module_op) == module_op_str_before +@pytest.mark.unit +def test_internal_state_reset(): + """Tests that the internal dag_builder state is reset.""" + pass + + @pytest.mark.unit class TestFuncOpVisualization: """Tests the visualization of FuncOps with bounding boxes""" From 805d21b0fe31ef3973d7f78a7b0b0ed279551dc7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:38:00 -0500 Subject: [PATCH 139/245] Revert "add test" This reverts commit 57ce573ea2369e60c2c5eff576b6c28bb3f8e370. --- .../visualization/test_construct_circuit_dag.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index ce8b9a1b02..55e6116324 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -128,12 +128,6 @@ def test_does_not_mutate_module(): assert str(module_op) == module_op_str_before -@pytest.mark.unit -def test_internal_state_reset(): - """Tests that the internal dag_builder state is reset.""" - pass - - @pytest.mark.unit class TestFuncOpVisualization: """Tests the visualization of FuncOps with bounding boxes""" From c0e0283cac27dfc7e35b3775e131137b33dbc521 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:38:17 -0500 Subject: [PATCH 140/245] Revert "try moving cluster stack to dag builder" This reverts commit 9847b3948d1012fdd1d0a25e20b8cd6c6bc1d062. --- .../visualization/construct_circuit_dag.py | 13 +++++++-- .../visualization/pydot_dag_builder.py | 28 ++----------------- .../test_construct_circuit_dag.py | 2 ++ 3 files changed, 14 insertions(+), 29 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 924132c3da..074bd50f03 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -41,9 +41,13 @@ class ConstructCircuitDAG: def __init__(self, dag_builder: DAGBuilder) -> None: self.dag_builder: DAGBuilder = dag_builder + # Keep track of nesting clusters using a stack + # NOTE: `None` corresponds to the base graph 'cluster' + self._cluster_stack: list[str | None] = [None] + def _reset(self) -> None: """Resets the instance.""" - self.dag_builder.reset() + self._cluster_stack: list[str | None] = [None] def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. @@ -87,6 +91,7 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: self.dag_builder.add_node( node_id, label=operation.device_name.data, + cluster_id=self._cluster_stack[-1], fillcolor="grey", color="black", penwidth=2, @@ -108,7 +113,9 @@ def _func_op(self, operation: func.FuncOp) -> None: self.dag_builder.add_cluster( cluster_id, label=operation.sym_name.data, + cluster_id=self._cluster_stack[-1], ) + self._cluster_stack.append(cluster_id) for region in operation.regions: self._visit_region(region) @@ -118,10 +125,10 @@ def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp - if len(self.dag_builder.cluster_id_stack) > 2: + if len(self._cluster_stack) > 2: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. - self.dag_builder.cluster_id_stack.pop() + self._cluster_stack.pop() for region in operation.regions: self._visit_region(region) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 5191ab23ba..5fd1cf7d74 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -65,9 +65,6 @@ def __init__( self._edges: list[dict[str, Any]] = [] self._clusters: dict[str, dict[str, Any]] = {} - # Keep track of nesting clusters using an internal stack - self.cluster_id_stack: list[str] = [] - _default_attrs: dict = ( {"fontname": "Helvetica", "penwidth": 2} if attrs is None else attrs ) @@ -101,20 +98,6 @@ def __init__( else cluster_attrs ) - def reset(self) -> None: - """Resets the instance.""" - self.graph: Dot = Dot( - graph_type="digraph", rankdir="TB", compound="true", strict=True - ) - - self._subgraph_cache: dict[str, Graph] = {} - - self._nodes: dict[str, dict[str, Any]] = {} - self._edges: list[dict[str, Any]] = [] - self._clusters: dict[str, dict[str, Any]] = {} - - self.cluster_id_stack: list[str] = [] - def add_node( self, id: str, @@ -217,19 +200,12 @@ def add_cluster( # Record new cluster self._subgraph_cache[id] = cluster + # Add node to cluster if cluster_id is None: - if self.cluster_id_stack: - # Nest on top of top most cluster - self._subgraph_cache[self.cluster_id_stack[-1]].add_subgraph(cluster) - else: - # If stack is empty, place on base graph - self.graph.add_subgraph(cluster) + self.graph.add_subgraph(cluster) else: self._subgraph_cache[cluster_id].add_subgraph(cluster) - # Add cluster to stack - self.cluster_id_stack.append(id) - self._clusters[id] = { "id": id, "cluster_label": cluster_attrs.get("label"), diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 55e6116324..a210a104ef 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,6 +16,8 @@ from unittest.mock import Mock import pytest +from _typeshed import GenericPath +from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") From 2a960cf387032a3e89f0fdeadcf89a098c159752 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:48:27 -0500 Subject: [PATCH 141/245] improve device test --- .../test_construct_circuit_dag.py | 90 +++++++++++++++++-- 1 file changed, 84 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a210a104ef..40af0fa136 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -24,15 +24,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -158,7 +157,9 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + generated_cluster_labels = { + info["cluster_label"] for info in graph_clusters.values() + } for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -270,5 +271,82 @@ def my_workflow(): for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) - cluster_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) + cluster_label_to_parent_label: dict[str, str] = dict( + zip(tuple(node_labels), parent_labels) + ) assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" + + def test_nested_qnodes(self): + """Tests that nested QJIT'd QNodes are visualized correctly""" + + dev1 = qml.device("null.qubit", wires=1) + dev2 = qml.device("lightning.qubit", wires=1) + + @qml.qnode(dev2) + def my_qnode2(): + qml.X(0) + + @qml.qnode(dev1) + def my_qnode1(): + qml.H(0) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + def my_workflow(): + my_qnode1() + my_qnode2() + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + graph_clusters = utility.dag_builder.clusters + graph_nodes = utility.dag_builder.nodes + + # Check labels we expected are there as clusters + expected_cluster_labels = [ + "jit_my_workflow", + "my_qnode1", + "my_qnode2", + "setup", + "teardown", + ] + assert len(graph_clusters) == len(expected_cluster_labels) + cluster_labels = {info["label"] for info in graph_clusters.values()} + for expected_name in expected_cluster_labels: + assert expected_name in cluster_labels + + # Check nesting is correct + # graph + # └── jit_my_workflow + # ├── my_qnode1 + # │ └── node: NullQubit + # └── my_qnode2 + # └── node: LightningQubit + + # Get the parent labels for each cluster and ensure they are what we expect. + parent_labels = ( + graph_clusters[child_cluster["cluster_id"]]["label"] + for child_cluster in graph_clusters.values() + if child_cluster.get("cluster_id") is not None + ) + cluster_label_to_parent_label: dict[str, str] = dict( + zip(tuple(cluster_labels), parent_labels) + ) + assert cluster_label_to_parent_label["jit_my_workflow"] is None + assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" + assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" + + # Check nodes are in the correct clusters + parent_labels = ( + graph_clusters[child_node["cluster_id"]]["label"] + for child_node in graph_nodes.values() + if child_node.get("cluster_id") is not None + ) + node_labels = {info["label"] for info in graph_nodes.values()} + node_label_to_parent_label: dict[str, str] = dict( + zip(tuple(node_labels), parent_labels) + ) + assert node_label_to_parent_label["NullQubit"] == "my_qnode1" + assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" From a21e6ddb02d8b1dfe2953c4771ad7f91087b9bb7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 14:51:53 -0500 Subject: [PATCH 142/245] fix black isort --- .../test_construct_circuit_dag.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 40af0fa136..6673d79546 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -24,14 +24,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -157,9 +158,7 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = { - info["cluster_label"] for info in graph_clusters.values() - } + generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -271,9 +270,7 @@ def my_workflow(): for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) - cluster_label_to_parent_label: dict[str, str] = dict( - zip(tuple(node_labels), parent_labels) - ) + cluster_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" def test_nested_qnodes(self): @@ -345,8 +342,6 @@ def my_workflow(): if child_node.get("cluster_id") is not None ) node_labels = {info["label"] for info in graph_nodes.values()} - node_label_to_parent_label: dict[str, str] = dict( - zip(tuple(node_labels), parent_labels) - ) + node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert node_label_to_parent_label["NullQubit"] == "my_qnode1" assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" From f666a9fe8876d4feac6e3a95414eb24c19ce04ab Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 15:20:10 -0500 Subject: [PATCH 143/245] add exceptions check --- .../visualization/pydot_dag_builder.py | 13 +++++ .../visualization/test_pydot_dag_builder.py | 49 ++++++++++++++++++- 2 files changed, 61 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index 5fd1cf7d74..e14241d073 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -114,6 +114,9 @@ def add_node( **attrs (Any): Any additional styling keyword arguments. """ + if id in self.nodes: + raise ValueError(f"Node ID {id} already present in graph.") + # Use ChainMap so you don't need to construct a new dictionary node_attrs: ChainMap = ChainMap(attrs, self._default_node_attrs) node = Node(id, label=label, **node_attrs) @@ -140,6 +143,13 @@ def add_edge(self, from_id: str, to_id: str, **attrs: Any) -> None: **attrs (Any): Any additional styling keyword arguments. """ + if from_id == to_id: + raise ValueError("Edges must connect two unique IDs.") + if from_id not in self.nodes: + raise ValueError("Source is not found in the graph.") + if to_id not in self.nodes: + raise ValueError("Destination is not found in the graph.") + # Use ChainMap so you don't need to construct a new dictionary edge_attrs: ChainMap = ChainMap(attrs, self._default_edge_attrs) edge = Edge(from_id, to_id, **edge_attrs) @@ -169,6 +179,9 @@ def add_cluster( **attrs (Any): Any additional styling keyword arguments. """ + if id in self.clusters: + raise ValueError(f"Cluster ID {id} already present in graph.") + # Use ChainMap so you don't need to construct a new dictionary cluster_attrs: ChainMap = ChainMap(attrs, self._default_cluster_attrs) cluster = Cluster(id, **cluster_attrs) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index b87b283903..5219889655 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -40,6 +40,53 @@ def test_initialization_defaults(): assert dag_builder.graph.obj_dict["strict"] is True +class TestExceptions: + """Tests the various exceptions defined in the class.""" + + def test_double_node_id(self): + """Tests that a ValueError is raised for duplicate nodes.""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + with pytest.raises(ValueError, match="Node ID 0 already present in graph."): + dag_builder.add_node("0", "node1") + + def test_edge_duplicate_source_destination(self): + """Tests that a ValueError is raised when an edge is created with the + same source and destination""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + with pytest.raises(ValueError, match="Edges must connect two uniques IDs."): + dag_builder.add_edge("0", "0") + + def test_edge_missing_ids(self): + """Tests that an error is raised if IDs are missing.""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("0", "node0") + with pytest.raises(ValueError, match="Destination is not found in the graph."): + dag_builder.add_edge("0", "1") + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_node("1", "node1") + with pytest.raises(ValueError, match="Source is not found in the graph."): + dag_builder.add_edge("0", "1") + + def test_duplicate_cluster_id(self): + """Tests that an exception is raised if an ID is already present.""" + + dag_builder = PyDotDAGBuilder() + + dag_builder.add_cluster("0") + with pytest.raises(ValueError, match="Cluster ID 0 already present in graph."): + dag_builder.add_cluster("0") + + class TestAddMethods: """Test that elements can be added to the graph.""" @@ -118,7 +165,7 @@ def test_add_cluster_to_parent_graph(self): # Level 0 (Root): Adds cluster on top of base graph dag_builder.add_node("n_root", "node_root") - # Level 1 (c0): Add node on outer cluster + # Level 1 (c0): Add node on outer cluster dag_builder.add_cluster("c0") dag_builder.add_node("n_outer", "node_outer", cluster_id="c0") From e28b2b71acbdc4aa3b7ae6cf402bdf9b43348526 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 15:23:43 -0500 Subject: [PATCH 144/245] add better documentation --- .../visualization/pydot_dag_builder.py | 10 ++++++++++ 1 file changed, 10 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py index e14241d073..ea01ecdf21 100644 --- a/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py +++ b/frontend/catalyst/python_interface/visualization/pydot_dag_builder.py @@ -113,6 +113,9 @@ def add_node( cluster_id (str | None): Optional ID of the cluster this node belongs to. **attrs (Any): Any additional styling keyword arguments. + Raises: + ValueError: Node ID is already present in the graph. + """ if id in self.nodes: raise ValueError(f"Node ID {id} already present in graph.") @@ -142,6 +145,11 @@ def add_edge(self, from_id: str, to_id: str, **attrs: Any) -> None: to_id (str): The unique ID of the destination node. **attrs (Any): Any additional styling keyword arguments. + Raises: + ValueError: Source and destination have the same ID + ValueError: Source is not found in the graph. + ValueError: Destination is not found in the graph. + """ if from_id == to_id: raise ValueError("Edges must connect two unique IDs.") @@ -178,6 +186,8 @@ def add_cluster( cluster_id (str | None): Optional ID of the cluster this cluster belongs to. If `None`, the cluster will be positioned on the base graph. **attrs (Any): Any additional styling keyword arguments. + Raises: + ValueError: Cluster ID is already present in the graph. """ if id in self.clusters: raise ValueError(f"Cluster ID {id} already present in graph.") From af5e52c7f45420b28a0c95f3160b2363734dd2a5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 15:27:01 -0500 Subject: [PATCH 145/245] fix --- .../visualization/test_construct_circuit_dag.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 6673d79546..5f1b34604c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,8 +16,6 @@ from unittest.mock import Mock import pytest -from _typeshed import GenericPath -from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") From 952fd7f21c05911bd1398e4c5935724d138b2501 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 15:36:35 -0500 Subject: [PATCH 146/245] fix typo --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 5219889655..197be4af70 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -59,7 +59,7 @@ def test_edge_duplicate_source_destination(self): dag_builder = PyDotDAGBuilder() dag_builder.add_node("0", "node0") - with pytest.raises(ValueError, match="Edges must connect two uniques IDs."): + with pytest.raises(ValueError, match="Edges must connect two unique IDs."): dag_builder.add_edge("0", "0") def test_edge_missing_ids(self): From 8f2dc98e0f987edef577902d7a3bd7303fb23399 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 15:38:59 -0500 Subject: [PATCH 147/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_pydot_dag_builder.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py index 197be4af70..d9a17f5e3b 100644 --- a/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py +++ b/frontend/test/pytest/python_interface/visualization/test_pydot_dag_builder.py @@ -43,7 +43,7 @@ def test_initialization_defaults(): class TestExceptions: """Tests the various exceptions defined in the class.""" - def test_double_node_id(self): + def test_duplicate_node_ids(self): """Tests that a ValueError is raised for duplicate nodes.""" dag_builder = PyDotDAGBuilder() From 7d042492e2a4177a57c370b2a86f4d9ada81a763 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:00:54 -0500 Subject: [PATCH 148/245] fix tests --- .../visualization/test_construct_circuit_dag.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 5f1b34604c..9c51749bbd 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -167,7 +167,7 @@ def my_workflow(): # Get the parent labels for each cluster and ensure they are what we expect. parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["label"] + graph_clusters[child_cluster["cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) @@ -224,7 +224,7 @@ def my_workflow(): # Get the parent labels for each cluster and ensure they are what we expect. parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["label"] + graph_clusters[child_cluster["cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) @@ -264,7 +264,7 @@ def my_workflow(): # Ensure nesting is correct graph_clusters = utility.dag_builder.clusters parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["label"] + graph_clusters[child_cluster["cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) @@ -322,7 +322,7 @@ def my_workflow(): # Get the parent labels for each cluster and ensure they are what we expect. parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["label"] + graph_clusters[child_cluster["cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) From 2bad538d86966c7ea0e95131f6f8bf9d125f3db3 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 12:11:33 -0500 Subject: [PATCH 149/245] add reminder comments --- .../visualization/test_construct_circuit_dag.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 0e14bd849d..ae2fe17be4 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -259,6 +258,8 @@ def my_workflow(): cluster_labels = {info["label"] for info in clusters.values()} assert "for ..." in cluster_labels + # Ensure proper nesting of clusters + @pytest.mark.unit def test_while_loop(self): """Test that the while loop is visualized correctly.""" @@ -282,6 +283,8 @@ def my_workflow(): cluster_labels = {info["label"] for info in clusters.values()} assert "while ..." in cluster_labels + # Ensure proper nesting of clusters + @pytest.mark.unit def test_if_else_conditional(self): """Test that the conditional operation is visualized correctly.""" @@ -307,6 +310,8 @@ def my_workflow(): assert "if ..." in cluster_labels assert "else" in cluster_labels + # Ensure proper nesting of clusters + @pytest.mark.unit def test_if_elif_else_conditional(self): """Test that the conditional operation is visualized correctly.""" @@ -336,7 +341,6 @@ def my_workflow(): assert "else" in cluster_labels assert cluster_labels.count("else") == 2 - class TestDeviceNode: """Tests that the device node is correctly visualized.""" From 683768e8534841100fe373751fe8915be0f11001 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:04:47 -0500 Subject: [PATCH 150/245] add back comment --- .../visualization/test_construct_circuit_dag.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index ae2fe17be4..edc9356760 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -341,6 +341,8 @@ def my_workflow(): assert "else" in cluster_labels assert cluster_labels.count("else") == 2 + # Ensure proper nesting + class TestDeviceNode: """Tests that the device node is correctly visualized.""" From 2adb3b2ea0fe2925be472df26e2c145f31ae30d8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:06:40 -0500 Subject: [PATCH 151/245] format --- .../visualization/test_construct_circuit_dag.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index edc9356760..f350facc3f 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -233,6 +234,7 @@ def my_workflow(): assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" + class TestControlFlowClusterVisualization: """Tests that the control flow operations are visualized correctly as clusters.""" @@ -343,6 +345,7 @@ def my_workflow(): # Ensure proper nesting + class TestDeviceNode: """Tests that the device node is correctly visualized.""" @@ -356,12 +359,12 @@ def test_standard_qnode(self): @qml.qnode(dev) def my_workflow(): qml.H(0) - + module = my_workflow() - + utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - + graph_nodes = utility.dag_builder.nodes # Basic check for node From f7bbb6bcefb1431bbcaefa5f6e9506948f02922b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:07:27 -0500 Subject: [PATCH 152/245] fix test name --- .../visualization/test_construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index f350facc3f..c4c4ceb8b9 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -235,7 +235,7 @@ def my_workflow(): assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" -class TestControlFlowClusterVisualization: +class TestControlFlowVisualization: """Tests that the control flow operations are visualized correctly as clusters.""" @pytest.mark.unit From f7f4b5ab45b3cd1c251042fa2573a262dbb2a47d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:16:32 -0500 Subject: [PATCH 153/245] fix test --- .../visualization/test_construct_circuit_dag.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 9c51749bbd..90dd3b9b54 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -212,7 +212,7 @@ def my_workflow(): "teardown", ] assert len(graph_clusters) == len(expected_cluster_labels) - cluster_labels = {info["label"] for info in graph_clusters.values()} + cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} for expected_name in expected_cluster_labels: assert expected_name in cluster_labels @@ -308,7 +308,7 @@ def my_workflow(): "teardown", ] assert len(graph_clusters) == len(expected_cluster_labels) - cluster_labels = {info["label"] for info in graph_clusters.values()} + cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} for expected_name in expected_cluster_labels: assert expected_name in cluster_labels @@ -335,7 +335,7 @@ def my_workflow(): # Check nodes are in the correct clusters parent_labels = ( - graph_clusters[child_node["cluster_id"]]["label"] + graph_clusters[child_node["cluster_id"]]["cluster_label"] for child_node in graph_nodes.values() if child_node.get("cluster_id") is not None ) From a21f878e98b2c3e1bb8b4cc57cd6cc8745fa7acd Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:22:08 -0500 Subject: [PATCH 154/245] fix tests --- .../visualization/construct_circuit_dag.py | 1 + .../test_construct_circuit_dag.py | 36 ++++++++----------- 2 files changed, 15 insertions(+), 22 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 074bd50f03..3bcc5e0921 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -109,6 +109,7 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: def _func_op(self, operation: func.FuncOp) -> None: """Visit a FuncOp Operation.""" + # Visualize the FuncOp as a cluster with a label cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 90dd3b9b54..f9ef3917ed 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -156,7 +155,9 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + generated_cluster_labels = { + info["cluster_label"] for info in graph_clusters.values() + } for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -256,20 +257,22 @@ def my_workflow(): utility.construct(module) graph_nodes = utility.dag_builder.nodes + graph_clusters = utility.dag_builder.clusters # Basic check for node node_labels = {info["label"] for info in graph_nodes.values()} assert "NullQubit" in node_labels # Ensure nesting is correct - graph_clusters = utility.dag_builder.clusters parent_labels = ( graph_clusters[child_cluster["cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) - cluster_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) - assert cluster_label_to_parent_label["NullQubit"] == "my_workflow" + node_label_to_parent_label: dict[str, str] = dict( + zip(tuple(node_labels), parent_labels) + ) + assert node_label_to_parent_label["NullQubit"] == "my_workflow" def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -320,19 +323,6 @@ def my_workflow(): # └── my_qnode2 # └── node: LightningQubit - # Get the parent labels for each cluster and ensure they are what we expect. - parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["cluster_label"] - for child_cluster in graph_clusters.values() - if child_cluster.get("cluster_id") is not None - ) - cluster_label_to_parent_label: dict[str, str] = dict( - zip(tuple(cluster_labels), parent_labels) - ) - assert cluster_label_to_parent_label["jit_my_workflow"] is None - assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" - assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" - # Check nodes are in the correct clusters parent_labels = ( graph_clusters[child_node["cluster_id"]]["cluster_label"] @@ -340,6 +330,8 @@ def my_workflow(): if child_node.get("cluster_id") is not None ) node_labels = {info["label"] for info in graph_nodes.values()} - node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) + node_label_to_parent_label: dict[str, str] = dict( + zip(tuple(node_labels), parent_labels) + ) assert node_label_to_parent_label["NullQubit"] == "my_qnode1" assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" From ae0ff325ad8063bf837073179d36f1687087afb1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:22:59 -0500 Subject: [PATCH 155/245] format --- .../test_construct_circuit_dag.py | 19 +++++++------------ 1 file changed, 7 insertions(+), 12 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index f9ef3917ed..177349df27 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -155,9 +156,7 @@ def my_workflow(): "setup", "teardown", ] - generated_cluster_labels = { - info["cluster_label"] for info in graph_clusters.values() - } + generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} for cluster_label in expected_cluster_labels: assert cluster_label in generated_cluster_labels @@ -269,9 +268,7 @@ def my_workflow(): for child_cluster in graph_clusters.values() if child_cluster.get("cluster_id") is not None ) - node_label_to_parent_label: dict[str, str] = dict( - zip(tuple(node_labels), parent_labels) - ) + node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert node_label_to_parent_label["NullQubit"] == "my_workflow" def test_nested_qnodes(self): @@ -330,8 +327,6 @@ def my_workflow(): if child_node.get("cluster_id") is not None ) node_labels = {info["label"] for info in graph_nodes.values()} - node_label_to_parent_label: dict[str, str] = dict( - zip(tuple(node_labels), parent_labels) - ) + node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert node_label_to_parent_label["NullQubit"] == "my_qnode1" assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" From 26ae96842da8e15191ab8bef8f57bf8095330102 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 26 Nov 2025 16:23:29 -0500 Subject: [PATCH 156/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 3bcc5e0921..02c5880e7f 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -42,7 +42,7 @@ def __init__(self, dag_builder: DAGBuilder) -> None: self.dag_builder: DAGBuilder = dag_builder # Keep track of nesting clusters using a stack - # NOTE: `None` corresponds to the base graph 'cluster' + # NOTE: `None` corresponds to the base graph self._cluster_stack: list[str | None] = [None] def _reset(self) -> None: From 1d4f1f1d10c9cdfbde38e407d3fd710fe496fee0 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:24:56 -0500 Subject: [PATCH 157/245] more explanation --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 02c5880e7f..188f74fcda 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -126,6 +126,7 @@ def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp + # and we want to use the jit_* workflow as the outer most bounding box. if len(self._cluster_stack) > 2: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. From ed8bae036f856f73c56144baa767795d0ee84a02 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 16:48:22 -0500 Subject: [PATCH 158/245] more fixes --- .../visualization/construct_circuit_dag.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 7c52a1d06d..e3583c8287 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -91,6 +91,7 @@ def _for_op(self, operation: scf.ForOp) -> None: self.dag_builder.add_cluster( cluster_id, node_label="for ...", + label="", cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) @@ -107,6 +108,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: self.dag_builder.add_cluster( cluster_id, node_label="while ...", + label="", cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) @@ -122,7 +124,8 @@ def _if_op(self, operation: scf.IfOp): cluster_id = f"cluster_{id(operation)}" self.dag_builder.add_cluster( cluster_id, - node_label="if", + node_label="", + label="conditional", cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) @@ -132,7 +135,8 @@ def _if_op(self, operation: scf.IfOp): cluster_id = f"cluster_ifop_branch{i}_{id(operation)}" self.dag_builder.add_cluster( cluster_id, - label=f"if ..." if i == 0 else "else", + node_label=f"if ..." if i == 0 else "else", + label="", cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) From 25e805105f112699b4ec7c5559bb6397330d8647 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 17:08:20 -0500 Subject: [PATCH 159/245] fix up testing --- .../test_construct_circuit_dag.py | 34 +++++++++++-------- 1 file changed, 20 insertions(+), 14 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 177349df27..b0e3205e9c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -49,7 +49,7 @@ def add_node(self, id, label, cluster_id=None, **attrs) -> None: self._nodes[id] = { "id": id, "label": label, - "cluster_id": cluster_id, + "parent_cluster_id": "base" if cluster_id is None else cluster_id, "attrs": attrs, } @@ -73,7 +73,7 @@ def add_cluster( "id": id, "node_label": node_label, "cluster_label": attrs.get("label"), - "cluster_id": cluster_id, + "parent_cluster_id": "base" if cluster_id is None else cluster_id, "attrs": attrs, } @@ -167,9 +167,8 @@ def my_workflow(): # Get the parent labels for each cluster and ensure they are what we expect. parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["cluster_label"] + graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() - if child_cluster.get("cluster_id") is not None ) cluster_label_to_parent_label: dict[str, str] = dict( zip(tuple(generated_cluster_labels), parent_labels) @@ -224,9 +223,9 @@ def my_workflow(): # Get the parent labels for each cluster and ensure they are what we expect. parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["cluster_label"] + graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() - if child_cluster.get("cluster_id") is not None + if child_cluster.get("parent_cluster_id") is not None ) cluster_label_to_parent_label: dict[str, str] = dict( zip(tuple(cluster_labels), parent_labels) @@ -264,9 +263,9 @@ def my_workflow(): # Ensure nesting is correct parent_labels = ( - graph_clusters[child_cluster["cluster_id"]]["cluster_label"] + graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] for child_cluster in graph_clusters.values() - if child_cluster.get("cluster_id") is not None + if child_cluster.get("parent_cluster_id") is not None ) node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert node_label_to_parent_label["NullQubit"] == "my_workflow" @@ -321,12 +320,19 @@ def my_workflow(): # └── node: LightningQubit # Check nodes are in the correct clusters - parent_labels = ( - graph_clusters[child_node["cluster_id"]]["cluster_label"] - for child_node in graph_nodes.values() - if child_node.get("cluster_id") is not None + nesting_relationships = dict( + ( + child_data["cluster_label"], # Child label (e.g., 'my_workflow') + # If the parent ID is the literal "base", assign the conceptual label. + ( + "base" + if child_data["parent_cluster_id"] == "base" + # Otherwise, perform the standard dictionary lookup using the parent ID. + else clusters[child_data["parent_cluster_id"]]["cluster_label"] + ), + ) + # Iterate over all cluster objects (the children) + for child_data in clusters.values() ) - node_labels = {info["label"] for info in graph_nodes.values()} - node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) assert node_label_to_parent_label["NullQubit"] == "my_qnode1" assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" From ffc9726bef6decc285fc10b7e4b1f4abfebbd4e5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 26 Nov 2025 17:10:50 -0500 Subject: [PATCH 160/245] fix naming --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 4a98739c5a..14b7892e6b 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -47,7 +47,7 @@ def add_node(self, id, label, cluster_id=None, **attrs) -> None: self._nodes[id] = { "id": id, "label": label, - "cluster_id": cluster_id, + "parent_cluster_id": cluster_id, "attrs": attrs, } @@ -71,7 +71,7 @@ def add_cluster( "id": id, "node_label": node_label, "cluster_label": attrs.get("label"), - "cluster_id": cluster_id, + "parent_cluster_id": cluster_id, "attrs": attrs, } From e58891a5195724cf531d4abcd5e69c4fd7802927 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 11:29:02 -0500 Subject: [PATCH 161/245] new testing approach --- .../test_construct_circuit_dag.py | 182 ++++++++---------- 1 file changed, 81 insertions(+), 101 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index b0e3205e9c..8e0c91d8c1 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -77,6 +76,52 @@ def add_cluster( "attrs": attrs, } + def get_nodes_in_cluster(self, cluster_label: str) -> list[str]: + """ + Returns a list of node labels that are direct children of the given cluster. + """ + node_ids = [] + cluster_id = self.get_cluster_id_by_label(cluster_label) + for node_data in self._nodes.values(): + if node_data["parent_cluster_id"] == cluster_id: + node_ids.append(node_data["label"]) + return node_ids + + def get_child_clusters(self, parent_cluster_label: str) -> list[str]: + """ + Returns a list of cluster labels that are direct children of the given parent cluster. + """ + parent_cluster_id = self.get_cluster_id_by_label(parent_cluster_label) + cluster_labels = [] + for cluster_data in self._clusters.values(): + if cluster_data["parent_cluster_id"] == parent_cluster_id: + cluster_label = ( + cluster_data["cluster_label"] or cluster_data["node_label"] + ) + cluster_labels.append(cluster_label) + return cluster_labels + + def get_node_id_by_label(self, label: str) -> str | None: + """ + Finds the ID of a node given its label. + Assumes labels are unique for testing purposes. + """ + for id, node_data in self._nodes.items(): + if node_data["label"] == label: + return id + return None + + def get_cluster_id_by_label(self, label: str) -> str | None: + """ + Finds the ID of a cluster given its label. + Assumes cluster labels are unique for testing purposes. + """ + for id, cluster_data in self._clusters.items(): + cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] + if cluster_label == label: + return id + return None + @property def nodes(self): return self._nodes @@ -147,34 +192,26 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_clusters = utility.dag_builder.clusters - # Check labels we expected are there - expected_cluster_labels = [ - "jit_my_workflow", - "my_workflow", - "setup", - "teardown", - ] - generated_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - for cluster_label in expected_cluster_labels: - assert cluster_label in generated_cluster_labels + graph_clusters = utility.dag_builder.clusters + all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + assert "jit_my_workflow" in all_cluster_labels + assert "my_workflow" in all_cluster_labels # Check nesting is correct # graph # └── jit_my_workflow # └── my_workflow - # Get the parent labels for each cluster and ensure they are what we expect. - parent_labels = ( - graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] - for child_cluster in graph_clusters.values() + # Check my_workflow is nested under jit_my_workflow + assert "my_workflow" in utility.dag_builder.get_child_clusters( + "jit_my_workflow" ) - cluster_label_to_parent_label: dict[str, str] = dict( - zip(tuple(generated_cluster_labels), parent_labels) + # Check that jit_my_workflow is the first cluster on top of the graph + jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label( + "jit_my_workflow" ) - assert cluster_label_to_parent_label["jit_my_workflow"] is None - assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" + assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -203,17 +240,11 @@ def my_workflow(): graph_clusters = utility.dag_builder.clusters # Check labels we expected are there as clusters - expected_cluster_labels = [ - "jit_my_workflow", - "my_qnode1", - "my_qnode2", - "setup", - "teardown", - ] - assert len(graph_clusters) == len(expected_cluster_labels) - cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - for expected_name in expected_cluster_labels: - assert expected_name in cluster_labels + graph_clusters = utility.dag_builder.clusters + all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + assert "jit_my_workflow" in all_cluster_labels + assert "my_qnode1" in all_cluster_labels + assert "my_qnode2" in all_cluster_labels # Check nesting is correct # graph @@ -221,18 +252,15 @@ def my_workflow(): # ├── my_qnode1 # └── my_qnode2 - # Get the parent labels for each cluster and ensure they are what we expect. - parent_labels = ( - graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] - for child_cluster in graph_clusters.values() - if child_cluster.get("parent_cluster_id") is not None - ) - cluster_label_to_parent_label: dict[str, str] = dict( - zip(tuple(cluster_labels), parent_labels) + # Check jit_my_workflow is under graph + jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label( + "jit_my_workflow" ) - assert cluster_label_to_parent_label["jit_my_workflow"] is None - assert cluster_label_to_parent_label["my_qnode1"] == "jit_my_workflow" - assert cluster_label_to_parent_label["my_qnode2"] == "jit_my_workflow" + assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" + + # Check both qnodes are under jit_my_workflow + assert "my_qnode1" in utility.dag_builder.get_child_clusters("jit_my_workflow") + assert "my_qnode2" in utility.dag_builder.get_child_clusters("jit_my_workflow") class TestDeviceNode: @@ -254,21 +282,9 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_nodes = utility.dag_builder.nodes - graph_clusters = utility.dag_builder.clusters - - # Basic check for node - node_labels = {info["label"] for info in graph_nodes.values()} - assert "NullQubit" in node_labels - - # Ensure nesting is correct - parent_labels = ( - graph_clusters[child_cluster["parent_cluster_id"]]["cluster_label"] - for child_cluster in graph_clusters.values() - if child_cluster.get("parent_cluster_id") is not None - ) - node_label_to_parent_label: dict[str, str] = dict(zip(tuple(node_labels), parent_labels)) - assert node_label_to_parent_label["NullQubit"] == "my_workflow" + # Check that device node is within the my_workflow cluster + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_workflow") + assert "NullQubit" in nodes_in_my_workflow def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -295,44 +311,8 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_clusters = utility.dag_builder.clusters - graph_nodes = utility.dag_builder.nodes - - # Check labels we expected are there as clusters - expected_cluster_labels = [ - "jit_my_workflow", - "my_qnode1", - "my_qnode2", - "setup", - "teardown", - ] - assert len(graph_clusters) == len(expected_cluster_labels) - cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - for expected_name in expected_cluster_labels: - assert expected_name in cluster_labels - - # Check nesting is correct - # graph - # └── jit_my_workflow - # ├── my_qnode1 - # │ └── node: NullQubit - # └── my_qnode2 - # └── node: LightningQubit - - # Check nodes are in the correct clusters - nesting_relationships = dict( - ( - child_data["cluster_label"], # Child label (e.g., 'my_workflow') - # If the parent ID is the literal "base", assign the conceptual label. - ( - "base" - if child_data["parent_cluster_id"] == "base" - # Otherwise, perform the standard dictionary lookup using the parent ID. - else clusters[child_data["parent_cluster_id"]]["cluster_label"] - ), - ) - # Iterate over all cluster objects (the children) - for child_data in clusters.values() - ) - assert node_label_to_parent_label["NullQubit"] == "my_qnode1" - assert node_label_to_parent_label["LightningSimulator"] == "my_qnode2" + # Check that device node is within the my_workflow cluster + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode1") + assert "NullQubit" in nodes_in_my_workflow + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode2") + assert "LightningSimulator" in nodes_in_my_workflow From 6d1ab993fff0019de118a23bdd1071fd6436f4eb Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 11:30:30 -0500 Subject: [PATCH 162/245] format --- .../test_construct_circuit_dag.py | 23 +++++++------------ 1 file changed, 8 insertions(+), 15 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 8e0c91d8c1..1e39bd1f7f 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -95,9 +96,7 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_id"] == parent_cluster_id: - cluster_label = ( - cluster_data["cluster_label"] or cluster_data["node_label"] - ) + cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] cluster_labels.append(cluster_label) return cluster_labels @@ -204,13 +203,9 @@ def my_workflow(): # └── my_workflow # Check my_workflow is nested under jit_my_workflow - assert "my_workflow" in utility.dag_builder.get_child_clusters( - "jit_my_workflow" - ) + assert "my_workflow" in utility.dag_builder.get_child_clusters("jit_my_workflow") # Check that jit_my_workflow is the first cluster on top of the graph - jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label( - "jit_my_workflow" - ) + jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label("jit_my_workflow") assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" def test_nested_qnodes(self): @@ -253,9 +248,7 @@ def my_workflow(): # └── my_qnode2 # Check jit_my_workflow is under graph - jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label( - "jit_my_workflow" - ) + jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label("jit_my_workflow") assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" # Check both qnodes are under jit_my_workflow From 710e6a89269a0c3f2cb0c8d1a08a9a2d923b8e65 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 12:15:03 -0500 Subject: [PATCH 163/245] test the fakebackend --- .../test_construct_circuit_dag.py | 62 +++++++++++++++++++ 1 file changed, 62 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 1e39bd1f7f..a37e6b6542 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -140,6 +140,68 @@ def to_string(self) -> str: return "graph" +class TestFakeDAGBuilder: + """Test the FakeDAGBuilder to ensure helper functions work as intended.""" + + @pytest.fixture + def builder_with_data(self): + """Sets up an instance with a complex graph already built.""" + + builder = FakeDAGBuilder() + + # Cluster set-up + builder.add_cluster("c0", label="Company") + builder.add_cluster("c1", label="Marketing", cluster_id="c0") + builder.add_cluster("c2", label="Finance", cluster_id="c0") + + # Node set-up + builder.add_node("n0", "CEO", cluster_id=None) # Add to base cluster "Company" + builder.add_node("n1", "Marketing Manager", cluster_id="c1") + builder.add_node("n2", "Finance Manager", cluster_id="c2") + + return builder + + # Test ID look up + + def test_get_node_id_by_label_success(self, builder_with_data): + assert builder_with_data.get_node_id_by_label("Finance Manager") == "n2" + assert builder_with_data.get_node_id_by_label("Marketing Manager") == "n1" + assert builder_with_data.get_node_id_by_label("CEO") == "n0" + + def test_get_node_id_by_label_failure(self, builder_with_data): + assert builder_with_data.get_node_id_by_label("Software Manager") is None + + def test_get_cluster_id_by_label_success(self, builder_with_data): + assert builder_with_data.get_cluster_id_by_label("Finance") == "c2" + assert builder_with_data.get_cluster_id_by_label("Marketing") == "c1" + assert builder_with_data.get_cluster_id_by_label("Company") == "c0" + + def test_get_cluster_id_by_label_failure(self, builder_with_data): + assert builder_with_data.get_cluster_id_by_label("Software") is None + + # Test relationship probing + + def test_node_heirarchy(self, builder_with_data): + finance_nodes = builder_with_data.get_nodes_in_cluster("Finance") + assert finance_nodes == ["Finance Manager"] + + marketing_nodes = builder_with_data.get_nodes_in_cluster("Marketing") + assert marketing_nodes == ["Marketing Manager"] + + company_nodes = builder_with_data.get_nodes_in_cluster("Company") + assert company_nodes == ["CEO"] + + def test_cluster_heirarchy(self, builder_with_data): + clusters_in_finance = builder_with_data.get_child_clusters("Finance") + assert not clusters_in_finance + + clusters_in_marketing = builder_with_data.get_child_clusters("Marketing") + assert not clusters_in_marketing + + clusters_in_company = builder_with_data.get_child_clusters("Company") + assert {"Finance", "Marketing"} == set(clusters_in_company) + + @pytest.mark.unit def test_dependency_injection(): """Tests that relevant dependencies are injected.""" From cb931d1f6baef1e7a7734a7b160d9326d7345e38 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 12:16:49 -0500 Subject: [PATCH 164/245] whoops forgot test --- .../visualization/test_construct_circuit_dag.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a37e6b6542..ff85203c09 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -201,6 +201,9 @@ def test_cluster_heirarchy(self, builder_with_data): clusters_in_company = builder_with_data.get_child_clusters("Company") assert {"Finance", "Marketing"} == set(clusters_in_company) + clusters_in_base = builder_with_data.get_child_clusters("base") + assert clusters_in_base == ["Company"] + @pytest.mark.unit def test_dependency_injection(): From 049a683378032b32f1cdd90ccda77df7958bc600 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 12:35:43 -0500 Subject: [PATCH 165/245] improve test --- .../visualization/test_construct_circuit_dag.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index ff85203c09..98eaf2591f 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -115,6 +115,9 @@ def get_cluster_id_by_label(self, label: str) -> str | None: Finds the ID of a cluster given its label. Assumes cluster labels are unique for testing purposes. """ + # Work around for base graph + if label == "base": + return "base" for id, cluster_data in self._clusters.items(): cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] if cluster_label == label: @@ -150,7 +153,7 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster("c0", label="Company") + builder.add_cluster("c0", label="Company", cluster_id=None) # Add to base graph builder.add_cluster("c1", label="Marketing", cluster_id="c0") builder.add_cluster("c2", label="Finance", cluster_id="c0") @@ -202,7 +205,7 @@ def test_cluster_heirarchy(self, builder_with_data): assert {"Finance", "Marketing"} == set(clusters_in_company) clusters_in_base = builder_with_data.get_child_clusters("base") - assert clusters_in_base == ["Company"] + assert clusters_in_base == ["Company"] @pytest.mark.unit From 9ea917e7368ba0a42f8df7cd3788ac5e9ef884fa Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 12:53:00 -0500 Subject: [PATCH 166/245] add new feature of jit_ prefix removal and single qnode --- .../visualization/construct_circuit_dag.py | 31 ++++++++++++++----- .../test_construct_circuit_dag.py | 28 +++++++---------- 2 files changed, 35 insertions(+), 24 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 188f74fcda..0c3d088bbc 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -19,7 +19,7 @@ from xdsl.dialects import builtin, func from xdsl.ir import Block, Operation, Region -from catalyst.python_interface.dialects import quantum +from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder @@ -109,14 +109,29 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: def _func_op(self, operation: func.FuncOp) -> None: """Visit a FuncOp Operation.""" + # If this is the jit_* FuncOp, only draw if there's more than one qnode (launch kernel) + # This avoids redundant nested clusters: jit_my_circuit -> my_circuit -> ... + visualize = True + label = operation.sym_name.data + if "jit_" in operation.sym_name.data: + num_qnodes = 0 + for op in operation.walk(): + if isinstance(op, catalyst.LaunchKernelOp): + num_qnodes += 1 + # Get everything after the jit_* prefix + label = str(label).split("_", maxsplit=1)[-1] + if num_qnodes == 1: + visualize = False + # Visualize the FuncOp as a cluster with a label - cluster_id = f"cluster_{id(operation)}" - self.dag_builder.add_cluster( - cluster_id, - label=operation.sym_name.data, - cluster_id=self._cluster_stack[-1], - ) - self._cluster_stack.append(cluster_id) + if visualize: + cluster_id = f"cluster_{id(operation)}" + self.dag_builder.add_cluster( + cluster_id, + label=label, + cluster_id=self._cluster_stack[-1], + ) + self._cluster_stack.append(cluster_id) for region in operation.regions: self._visit_region(region) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 98eaf2591f..e8e66883bc 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -262,19 +262,15 @@ def my_workflow(): # Check labels we expected are there graph_clusters = utility.dag_builder.clusters all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - assert "jit_my_workflow" in all_cluster_labels assert "my_workflow" in all_cluster_labels # Check nesting is correct # graph - # └── jit_my_workflow - # └── my_workflow + # └── my_workflow - # Check my_workflow is nested under jit_my_workflow - assert "my_workflow" in utility.dag_builder.get_child_clusters("jit_my_workflow") - # Check that jit_my_workflow is the first cluster on top of the graph - jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label("jit_my_workflow") - assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" + # Check my_workflow is nested under my_workflow + my_workflow_id = utility.dag_builder.get_cluster_id_by_label("my_workflow") + assert graph_clusters[my_workflow_id]["parent_cluster_id"] == "base" def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -305,23 +301,23 @@ def my_workflow(): # Check labels we expected are there as clusters graph_clusters = utility.dag_builder.clusters all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - assert "jit_my_workflow" in all_cluster_labels + assert "my_workflow" in all_cluster_labels assert "my_qnode1" in all_cluster_labels assert "my_qnode2" in all_cluster_labels # Check nesting is correct # graph - # └── jit_my_workflow + # └── my_workflow # ├── my_qnode1 # └── my_qnode2 - # Check jit_my_workflow is under graph - jit_my_workflow_id = utility.dag_builder.get_cluster_id_by_label("jit_my_workflow") - assert graph_clusters[jit_my_workflow_id]["parent_cluster_id"] == "base" + # Check my_workflow is under graph + my_workflow_id = utility.dag_builder.get_cluster_id_by_label("my_workflow") + assert graph_clusters[my_workflow_id]["parent_cluster_id"] == "base" - # Check both qnodes are under jit_my_workflow - assert "my_qnode1" in utility.dag_builder.get_child_clusters("jit_my_workflow") - assert "my_qnode2" in utility.dag_builder.get_child_clusters("jit_my_workflow") + # Check both qnodes are under my_workflow + assert "my_qnode1" in utility.dag_builder.get_child_clusters("my_workflow") + assert "my_qnode2" in utility.dag_builder.get_child_clusters("my_workflow") class TestDeviceNode: From fd26a37816b385c6dee7f06e90d9cb5fa394e543 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 13:06:53 -0500 Subject: [PATCH 167/245] whoops --- .../visualization/test_construct_circuit_dag.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index e8e66883bc..cf3b76d4ce 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -96,7 +95,9 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_id"] == parent_cluster_id: - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] + cluster_label = ( + cluster_data["cluster_label"] or cluster_data["node_label"] + ) cluster_labels.append(cluster_label) return cluster_labels @@ -158,7 +159,7 @@ def builder_with_data(self): builder.add_cluster("c2", label="Finance", cluster_id="c0") # Node set-up - builder.add_node("n0", "CEO", cluster_id=None) # Add to base cluster "Company" + builder.add_node("n0", "CEO", cluster_id="c0") builder.add_node("n1", "Marketing Manager", cluster_id="c1") builder.add_node("n2", "Finance Manager", cluster_id="c2") From 1ad2ec0c3ca0a5ac63cb31e737682ee576daca23 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 13:07:02 -0500 Subject: [PATCH 168/245] format --- .../visualization/test_construct_circuit_dag.py | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index cf3b76d4ce..0b60c887db 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -95,9 +96,7 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_id"] == parent_cluster_id: - cluster_label = ( - cluster_data["cluster_label"] or cluster_data["node_label"] - ) + cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] cluster_labels.append(cluster_label) return cluster_labels From 8fc4a303770a6627ec6fbb2fc0521260107300e2 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 13:41:10 -0500 Subject: [PATCH 169/245] improve some tests --- .../test_construct_circuit_dag.py | 41 +++++-------------- 1 file changed, 11 insertions(+), 30 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 58a108c353..e45c83ed50 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,15 +22,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -96,7 +95,9 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_id"] == parent_cluster_id: - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] + cluster_label = ( + cluster_data["cluster_label"] or cluster_data["node_label"] + ) cluster_labels.append(cluster_label) return cluster_labels @@ -341,11 +342,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.clusters - cluster_labels = {info["label"] for info in clusters.values()} - assert "for ..." in cluster_labels - - # Ensure proper nesting of clusters + assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") @pytest.mark.unit def test_while_loop(self): @@ -366,11 +363,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.clusters - cluster_labels = {info["label"] for info in clusters.values()} - assert "while ..." in cluster_labels - - # Ensure proper nesting of clusters + assert "while ..." in utility.dag_builder.get_child_clusters("my_workflow") @pytest.mark.unit def test_if_else_conditional(self): @@ -392,12 +385,9 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.clusters - cluster_labels = {info["label"] for info in clusters.values()} - assert "if ..." in cluster_labels - assert "else" in cluster_labels - - # Ensure proper nesting of clusters + assert "conditional" in utility.dag_builder.get_child_clusters("my_workflow") + assert "if ..." in utility.dag_builder.get_child_clusters("conditional") + assert "else" in utility.dag_builder.get_child_clusters("conditional") @pytest.mark.unit def test_if_elif_else_conditional(self): @@ -421,15 +411,6 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - clusters = utility.dag_builder.clusters - cluster_labels = [info["label"] for info in clusters.values()] - assert "if ..." in cluster_labels - assert cluster_labels.count("if ...") == 2 - assert "else" in cluster_labels - assert cluster_labels.count("else") == 2 - - # Ensure proper nesting - class TestDeviceNode: """Tests that the device node is correctly visualized.""" From bfd28b247d910d7d6068446abb21c6ad8be8bb63 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 27 Nov 2025 18:14:30 -0500 Subject: [PATCH 170/245] shouldn't have conditional label --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 878f89dd89..a6eb191fd8 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -125,7 +125,7 @@ def _if_op(self, operation: scf.IfOp): self.dag_builder.add_cluster( cluster_id, node_label="", - label="conditional", + label="", cluster_id=self._cluster_stack[-1], ) self._cluster_stack.append(cluster_id) From 9aef39957dfa8c3d513f21b696583bd7dffb219c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 15:19:26 -0500 Subject: [PATCH 171/245] id to uid rename --- .../test_construct_circuit_dag.py | 62 +++++++++---------- 1 file changed, 31 insertions(+), 31 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index a435f4a678..f08a83e0af 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -84,26 +84,26 @@ def get_nodes_in_cluster(self, cluster_label: str) -> list[str]: """ Returns a list of node labels that are direct children of the given cluster. """ - node_ids = [] - cluster_id = self.get_cluster_id_by_label(cluster_label) + node_uids = [] + cluster_uid = self.get_cluster_uid_by_label(cluster_label) for node_data in self._nodes.values(): - if node_data["parent_cluster_id"] == cluster_id: - node_ids.append(node_data["label"]) - return node_ids + if node_data["parent_cluster_uid"] == cluster_uid: + node_uids.append(node_data["label"]) + return node_uids def get_child_clusters(self, parent_cluster_label: str) -> list[str]: """ Returns a list of cluster labels that are direct children of the given parent cluster. """ - parent_cluster_id = self.get_cluster_id_by_label(parent_cluster_label) + parent_cluster_uid = self.get_cluster_uid_by_label(parent_cluster_label) cluster_labels = [] for cluster_data in self._clusters.values(): - if cluster_data["parent_cluster_id"] == parent_cluster_id: + if cluster_data["parent_cluster_uid"] == parent_cluster_uid: cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] cluster_labels.append(cluster_label) return cluster_labels - def get_node_id_by_label(self, label: str) -> str | None: + def get_node_uid_by_label(self, label: str) -> str | None: """ Finds the ID of a node given its label. Assumes labels are unique for testing purposes. @@ -113,7 +113,7 @@ def get_node_id_by_label(self, label: str) -> str | None: return id return None - def get_cluster_id_by_label(self, label: str) -> str | None: + def get_cluster_uid_by_label(self, label: str) -> str | None: """ Finds the ID of a cluster given its label. Assumes cluster labels are unique for testing purposes. @@ -156,34 +156,34 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster("c0", label="Company", cluster_id=None) # Add to base graph - builder.add_cluster("c1", label="Marketing", cluster_id="c0") - builder.add_cluster("c2", label="Finance", cluster_id="c0") + builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph + builder.add_cluster("c1", label="Marketing", cluster_uid="c0") + builder.add_cluster("c2", label="Finance", cluster_uid="c0") # Node set-up - builder.add_node("n0", "CEO", cluster_id="c0") - builder.add_node("n1", "Marketing Manager", cluster_id="c1") - builder.add_node("n2", "Finance Manager", cluster_id="c2") + builder.add_node("n0", "CEO", cluster_uid="c0") + builder.add_node("n1", "Marketing Manager", cluster_uid="c1") + builder.add_node("n2", "Finance Manager", cluster_uid="c2") return builder # Test ID look up - def test_get_node_id_by_label_success(self, builder_with_data): - assert builder_with_data.get_node_id_by_label("Finance Manager") == "n2" - assert builder_with_data.get_node_id_by_label("Marketing Manager") == "n1" - assert builder_with_data.get_node_id_by_label("CEO") == "n0" + def test_get_node_uid_by_label_success(self, builder_with_data): + assert builder_with_data.get_node_uid_by_label("Finance Manager") == "n2" + assert builder_with_data.get_node_uid_by_label("Marketing Manager") == "n1" + assert builder_with_data.get_node_uid_by_label("CEO") == "n0" - def test_get_node_id_by_label_failure(self, builder_with_data): - assert builder_with_data.get_node_id_by_label("Software Manager") is None + def test_get_node_uid_by_label_failure(self, builder_with_data): + assert builder_with_data.get_node_uid_by_label("Software Manager") is None - def test_get_cluster_id_by_label_success(self, builder_with_data): - assert builder_with_data.get_cluster_id_by_label("Finance") == "c2" - assert builder_with_data.get_cluster_id_by_label("Marketing") == "c1" - assert builder_with_data.get_cluster_id_by_label("Company") == "c0" + def test_get_cluster_uid_by_label_success(self, builder_with_data): + assert builder_with_data.get_cluster_uid_by_label("Finance") == "c2" + assert builder_with_data.get_cluster_uid_by_label("Marketing") == "c1" + assert builder_with_data.get_cluster_uid_by_label("Company") == "c0" - def test_get_cluster_id_by_label_failure(self, builder_with_data): - assert builder_with_data.get_cluster_id_by_label("Software") is None + def test_get_cluster_uid_by_label_failure(self, builder_with_data): + assert builder_with_data.get_cluster_uid_by_label("Software") is None # Test relationship probing @@ -272,8 +272,8 @@ def my_workflow(): # └── my_workflow # Check my_workflow is nested under my_workflow - my_workflow_id = utility.dag_builder.get_cluster_id_by_label("my_workflow") - assert graph_clusters[my_workflow_id]["parent_cluster_id"] == "base" + my_workflow_id = utility.dag_builder.get_cluster_uid_by_label("my_workflow") + assert graph_clusters[my_workflow_id]["parent_cluster_uid"] == "base" def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -315,8 +315,8 @@ def my_workflow(): # └── my_qnode2 # Check my_workflow is under graph - my_workflow_id = utility.dag_builder.get_cluster_id_by_label("my_workflow") - assert graph_clusters[my_workflow_id]["parent_cluster_id"] == "base" + my_workflow_id = utility.dag_builder.get_cluster_uid_by_label("my_workflow") + assert graph_clusters[my_workflow_id]["parent_cluster_uid"] == "base" # Check both qnodes are under my_workflow assert "my_qnode1" in utility.dag_builder.get_child_clusters("my_workflow") From e87dba289c9e71c86f43d136b811e4013e9e7bb3 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 15:19:37 -0500 Subject: [PATCH 172/245] format --- .../visualization/test_construct_circuit_dag.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index f08a83e0af..093c2966ba 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -31,9 +31,6 @@ ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From be1961c5495810dd381400f73ca76f6f84b353d8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 15:35:15 -0500 Subject: [PATCH 173/245] rename inner stack --- .../visualization/construct_circuit_dag.py | 20 +++++++++---------- 1 file changed, 10 insertions(+), 10 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0c3d088bbc..00038bf659 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -42,12 +42,11 @@ def __init__(self, dag_builder: DAGBuilder) -> None: self.dag_builder: DAGBuilder = dag_builder # Keep track of nesting clusters using a stack - # NOTE: `None` corresponds to the base graph - self._cluster_stack: list[str | None] = [None] + self._cluster_uid_stack: list[str] = [] def _reset(self) -> None: """Resets the instance.""" - self._cluster_stack: list[str | None] = [None] + self._cluster_uid_stack: list[str] = [] def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. @@ -91,7 +90,7 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: self.dag_builder.add_node( node_id, label=operation.device_name.data, - cluster_id=self._cluster_stack[-1], + cluster_id=self._cluster_uid_stack[-1], fillcolor="grey", color="black", penwidth=2, @@ -129,9 +128,9 @@ def _func_op(self, operation: func.FuncOp) -> None: self.dag_builder.add_cluster( cluster_id, label=label, - cluster_id=self._cluster_stack[-1], + cluster_id=self._cluster_uid_stack[-1], ) - self._cluster_stack.append(cluster_id) + self._cluster_uid_stack.append(cluster_id) for region in operation.regions: self._visit_region(region) @@ -140,12 +139,13 @@ def _func_op(self, operation: func.FuncOp) -> None: def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" - # NOTE: Skip first two because the first is the base graph, second is the jit_* workflow FuncOp - # and we want to use the jit_* workflow as the outer most bounding box. - if len(self._cluster_stack) > 2: + # NOTE: Skip first cluster as it is the "base" of the graph diagram. + # If it is a multi-qnode workflow, it will represent the "workflow" function + # If it is a single qnode, it will represent the quantum function. + if len(self._cluster_uid_stack) > 1: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. - self._cluster_stack.pop() + self._cluster_uid_stack.pop() for region in operation.regions: self._visit_region(region) From c92736a15f460c3546f0710998a64d91f8350f60 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 15:35:33 -0500 Subject: [PATCH 174/245] format --- .../visualization/test_construct_circuit_dag.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 18637674bd..093c2966ba 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -19,10 +19,6 @@ pytestmark = pytest.mark.usefixtures("requires_xdsl") -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml From ca44e59fdbdb624bceacdf87f48c55e176df3ea0 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Fri, 28 Nov 2025 15:40:12 -0500 Subject: [PATCH 175/245] Apply suggestion from @andrijapau --- .../visualization/test_construct_circuit_dag.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 093c2966ba..051a31e690 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -212,9 +212,9 @@ def test_cluster_heirarchy(self, builder_with_data): def test_dependency_injection(): """Tests that relevant dependencies are injected.""" - dag_builder = FakeDAGBuilder() - utility = ConstructCircuitDAG(dag_builder) - assert utility.dag_builder is dag_builder + mock_dag_builder = Mock(DAGBuilder) + utility = ConstructCircuitDAG(mock_dag_builder) + assert utility.dag_builder is mock_dag_builder @pytest.mark.unit From 3a8da550cb3f3a5c999b648c6c6e2c6bd9e1da43 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Fri, 28 Nov 2025 16:04:12 -0500 Subject: [PATCH 176/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 00038bf659..914dc28629 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -125,10 +125,11 @@ def _func_op(self, operation: func.FuncOp) -> None: # Visualize the FuncOp as a cluster with a label if visualize: cluster_id = f"cluster_{id(operation)}" + parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] self.dag_builder.add_cluster( cluster_id, label=label, - cluster_id=self._cluster_uid_stack[-1], + cluster_id=parent_cluster_uid, ) self._cluster_uid_stack.append(cluster_id) From 26e95f9cb335e98aed3716bb49dbc2be7f23a06b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 16:04:26 -0500 Subject: [PATCH 177/245] format --- .../python_interface/visualization/construct_circuit_dag.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 914dc28629..b4dedc934d 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -125,7 +125,9 @@ def _func_op(self, operation: func.FuncOp) -> None: # Visualize the FuncOp as a cluster with a label if visualize: cluster_id = f"cluster_{id(operation)}" - parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] + parent_cluster_uid = ( + None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] + ) self.dag_builder.add_cluster( cluster_id, label=label, From c9594ee6c142ac16434a30fba7744cdbb3099aa5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 16:28:59 -0500 Subject: [PATCH 178/245] fix --- .../visualization/construct_circuit_dag.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index b4dedc934d..2d3265617e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -90,7 +90,7 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: self.dag_builder.add_node( node_id, label=operation.device_name.data, - cluster_id=self._cluster_uid_stack[-1], + cluster_uid=self._cluster_uid_stack[-1], fillcolor="grey", color="black", penwidth=2, @@ -124,16 +124,16 @@ def _func_op(self, operation: func.FuncOp) -> None: # Visualize the FuncOp as a cluster with a label if visualize: - cluster_id = f"cluster_{id(operation)}" + uid = f"cluster_{id(operation)}" parent_cluster_uid = ( None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] ) self.dag_builder.add_cluster( - cluster_id, + uid, label=label, - cluster_id=parent_cluster_uid, + cluster_uid=parent_cluster_uid, ) - self._cluster_uid_stack.append(cluster_id) + self._cluster_uid_stack.append(uid) for region in operation.regions: self._visit_region(region) From c0554c5c7f9e4b2161ab9ef2e411b36cb0ffd69f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 16:33:19 -0500 Subject: [PATCH 179/245] fix --- .../visualization/construct_circuit_dag.py | 32 +++++++++---------- 1 file changed, 16 insertions(+), 16 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index b74338c7fd..97e55655d8 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -86,14 +86,14 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" - cluster_id = f"cluster_{id(operation)}" + uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( - cluster_id, + uid, node_label="for ...", label="", - cluster_id=self._cluster_stack[-1], + cluster_uid=self._cluster_stack[-1], ) - self._cluster_stack.append(cluster_id) + self._cluster_stack.append(uid) for region in operation.regions: self._visit_region(region) @@ -103,14 +103,14 @@ def _for_op(self, operation: scf.ForOp) -> None: @_visit_operation.register def _while_op(self, operation: scf.WhileOp) -> None: """Handle an xDSL WhileOp operation.""" - cluster_id = f"cluster_{id(operation)}" + uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( - cluster_id, + uid, node_label="while ...", label="", - cluster_id=self._cluster_stack[-1], + cluster_uid=self._cluster_stack[-1], ) - self._cluster_stack.append(cluster_id) + self._cluster_stack.append(uid) for region in operation.regions: self._visit_region(region) @@ -120,25 +120,25 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - cluster_id = f"cluster_{id(operation)}" + uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( - cluster_id, + uid, node_label="", label="", - cluster_id=self._cluster_stack[-1], + cluster_uid=self._cluster_stack[-1], ) - self._cluster_stack.append(cluster_id) + self._cluster_stack.append(uid) # Loop through each branch and visualize as a cluster for i, branch in enumerate(operation.regions): - cluster_id = f"cluster_ifop_branch{i}_{id(operation)}" + uid = f"cluster_ifop_branch{i}_{id(operation)}" self.dag_builder.add_cluster( - cluster_id, + uid, node_label=f"if ..." if i == 0 else "else", label="", - cluster_id=self._cluster_stack[-1], + cluster_uid=self._cluster_stack[-1], ) - self._cluster_stack.append(cluster_id) + self._cluster_stack.append(uid) # Go recursively into the branch to process internals self._visit_region(branch) From a80141c469bc04d976924c19009f22907d4a6f34 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 16:34:01 -0500 Subject: [PATCH 180/245] format --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index ebc129be64..50a397dd86 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -22,14 +22,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): From 529cace1d3e25248e50bf9301b4637db4ece7e68 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 28 Nov 2025 16:35:19 -0500 Subject: [PATCH 181/245] rename cluster stack --- .../visualization/construct_circuit_dag.py | 24 +++++++++---------- 1 file changed, 12 insertions(+), 12 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 97e55655d8..c6302ccb8e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -91,14 +91,14 @@ def _for_op(self, operation: scf.ForOp) -> None: uid, node_label="for ...", label="", - cluster_uid=self._cluster_stack[-1], + cluster_uid=self._cluster_uid_stack[-1], ) - self._cluster_stack.append(uid) + self._cluster_uid_stack.append(uid) for region in operation.regions: self._visit_region(region) - self._cluster_stack.pop() + self._cluster_uid_stack.pop() @_visit_operation.register def _while_op(self, operation: scf.WhileOp) -> None: @@ -108,14 +108,14 @@ def _while_op(self, operation: scf.WhileOp) -> None: uid, node_label="while ...", label="", - cluster_uid=self._cluster_stack[-1], + cluster_uid=self._cluster_uid_stack[-1], ) - self._cluster_stack.append(uid) + self._cluster_uid_stack.append(uid) for region in operation.regions: self._visit_region(region) - self._cluster_stack.pop() + self._cluster_uid_stack.pop() @_visit_operation.register def _if_op(self, operation: scf.IfOp): @@ -125,9 +125,9 @@ def _if_op(self, operation: scf.IfOp): uid, node_label="", label="", - cluster_uid=self._cluster_stack[-1], + cluster_uid=self._cluster_uid_stack[-1], ) - self._cluster_stack.append(uid) + self._cluster_uid_stack.append(uid) # Loop through each branch and visualize as a cluster for i, branch in enumerate(operation.regions): @@ -136,18 +136,18 @@ def _if_op(self, operation: scf.IfOp): uid, node_label=f"if ..." if i == 0 else "else", label="", - cluster_uid=self._cluster_stack[-1], + cluster_uid=self._cluster_uid_stack[-1], ) - self._cluster_stack.append(uid) + self._cluster_uid_stack.append(uid) # Go recursively into the branch to process internals self._visit_region(branch) # Pop branch cluster after processing to ensure # logical branches are treated as 'parallel' - self._cluster_stack.pop() + self._cluster_uid_stack.pop() - self._cluster_stack.pop() + self._cluster_uid_stack.pop() # ============ # DEVICE NODE From 4a172dd9195db6830b295a6d299d8bfb0bbc253f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 11:26:05 -0500 Subject: [PATCH 182/245] return op doesnt have regions --- .../python_interface/visualization/construct_circuit_dag.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 2d3265617e..8e492166a6 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -149,6 +149,3 @@ def _func_return(self, operation: func.ReturnOp) -> None: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. self._cluster_uid_stack.pop() - - for region in operation.regions: - self._visit_region(region) From a54de217d7f1171aa283eb606020869e76911587 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Mon, 1 Dec 2025 11:26:51 -0500 Subject: [PATCH 183/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 8e492166a6..1aeb3ceee7 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -122,7 +122,6 @@ def _func_op(self, operation: func.FuncOp) -> None: if num_qnodes == 1: visualize = False - # Visualize the FuncOp as a cluster with a label if visualize: uid = f"cluster_{id(operation)}" parent_cluster_uid = ( From 2bb5bdeaf37734ae2795ad32a7c17e86a20ea1c5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:24:28 -0500 Subject: [PATCH 184/245] update testing and add for loop label upgrade --- .../visualization/construct_circuit_dag.py | 11 ++++++++- .../test_construct_circuit_dag.py | 23 ++++++++++++++----- 2 files changed, 27 insertions(+), 7 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index b0e6ff6714..2f90832d77 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -20,6 +20,7 @@ from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import catalyst, quantum +from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder @@ -86,10 +87,18 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" + lower_bound, upper_bound, step = ( + resolve_constant_params(operation.lb), + resolve_constant_params(operation.ub), + resolve_constant_params(operation.step), + ) + + index_var_name = operation.body.blocks[0].args[0].name_hint + uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( uid, - node_label="for ...", + node_label=f"for {index_var_name} in range({lower_bound},{upper_bound},{step})", label="", cluster_uid=self._cluster_uid_stack[-1], ) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 50a397dd86..9ba891f1a1 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -13,6 +13,7 @@ # limitations under the License. """Unit tests for the ConstructCircuitDAG utility.""" +import re from unittest.mock import Mock import pytest @@ -320,12 +321,11 @@ def my_workflow(): assert "my_qnode2" in utility.dag_builder.get_child_clusters("my_workflow") -class TestControlFlowVisualization: - """Tests that the control flow operations are visualized correctly as clusters.""" +class TestForOp: + """Tests that the for loop control flow can be visualized correctly.""" - @pytest.mark.unit - def test_for_loop(self): - """Test that the for loop is visualized correctly.""" + def test_cluster_visualization(self): + """Tests that the for loop cluster can be visualized correctly.""" dev = qml.device("null.qubit", wires=1) @@ -341,7 +341,14 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") + assert re.search( + r"for arg\d in range\(0,2,1\)", + utility.dag_builder.get_child_clusters("my_workflow"), + ) + + +class TestWhileOp: + """Tests that the while loop control flow can be visualized correctly.""" @pytest.mark.unit def test_while_loop(self): @@ -364,6 +371,10 @@ def my_workflow(): assert "while ..." in utility.dag_builder.get_child_clusters("my_workflow") + +class TestIfOp: + """Tests that the conditional control flow can be visualized correctly.""" + @pytest.mark.unit def test_if_else_conditional(self): """Test that the conditional operation is visualized correctly.""" From 434c50013df2e48869264301ca06b2c64e70ac83 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:25:58 -0500 Subject: [PATCH 185/245] fix --- .../visualization/test_construct_circuit_dag.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 9ba891f1a1..bb8af4783f 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -324,7 +324,8 @@ def my_workflow(): class TestForOp: """Tests that the for loop control flow can be visualized correctly.""" - def test_cluster_visualization(self): + @pytest.mark.unit + def test_basic_example(self): """Tests that the for loop cluster can be visualized correctly.""" dev = qml.device("null.qubit", wires=1) @@ -351,7 +352,7 @@ class TestWhileOp: """Tests that the while loop control flow can be visualized correctly.""" @pytest.mark.unit - def test_while_loop(self): + def test_basic_example(self): """Test that the while loop is visualized correctly.""" dev = qml.device("null.qubit", wires=1) @@ -376,7 +377,7 @@ class TestIfOp: """Tests that the conditional control flow can be visualized correctly.""" @pytest.mark.unit - def test_if_else_conditional(self): + def test_basic_example(self): """Test that the conditional operation is visualized correctly.""" dev = qml.device("null.qubit", wires=1) From cfa5f3b29715b493f6b282314b6bab17cea8dc3e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:33:27 -0500 Subject: [PATCH 186/245] add more tests --- .../test_construct_circuit_dag.py | 35 ++++++++++++++++++- 1 file changed, 34 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index bb8af4783f..205a7b03a4 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -343,10 +343,43 @@ def my_workflow(): utility.construct(module) assert re.search( - r"for arg\d in range\(0,2,1\)", + r"for arg\d in range\(0,3,1\)", utility.dag_builder.get_child_clusters("my_workflow"), ) + @pytest.mark.unit + def test_nested_loop(self): + """Tests that nested for loops are visualized correctly.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + for i in range(0, 5, 2): + for j in range(1, 6, 2): + qml.H(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + child_clusters = utility.dag_builder.get_child_clusters("my_workflow") + assert len(child_clusters) == 1 + assert re.search( + r"for arg\d in range\(0,5,2\)", + child_clusters, + ) + for_loop_label = utility.dag_builder.get_child_clusters(child_clusters[0]) + child_clusters = utility.dag_builder.get_child_clusters(for_loop_label) + assert len(child_clusters) == 1 + assert re.search( + r"for arg\d in range\(1,6,2\)", + child_clusters, + ) + class TestWhileOp: """Tests that the while loop control flow can be visualized correctly.""" From a62f2b1055b0ab2bf5ba32dd4bbbd0fa9bd5a1a2 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:35:43 -0500 Subject: [PATCH 187/245] add dev comments --- .../visualization/test_construct_circuit_dag.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 205a7b03a4..b45b6d82d6 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -366,12 +366,15 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) + # Check first for loop child_clusters = utility.dag_builder.get_child_clusters("my_workflow") assert len(child_clusters) == 1 assert re.search( r"for arg\d in range\(0,5,2\)", child_clusters, ) + + # Check second for loop for_loop_label = utility.dag_builder.get_child_clusters(child_clusters[0]) child_clusters = utility.dag_builder.get_child_clusters(for_loop_label) assert len(child_clusters) == 1 From d9026ef010aaaca1824af552681e042bf96a9202 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:46:12 -0500 Subject: [PATCH 188/245] more testing --- .../test_construct_circuit_dag.py | 50 +++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index b45b6d82d6..370e7de4e7 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -383,6 +383,56 @@ def my_workflow(): child_clusters, ) + def test_dynamic_start_stop_step(self): + """Tests that dynamic start, stop, step variables can be displayed.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def step_dynamic(x): + for i in range(0, 3, x): + qml.H(0) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def stop_dynamic(x): + for i in range(0, x, 1): + qml.H(0) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def start_dynamic(x): + for i in range(x, 3, 1): + qml.H(0) + + args = (1,) + start_dynamic_module = start_dynamic(*args) + stop_dynamic_module = stop_dynamic(*args) + step_dynamic_module = step_dynamic(*args) + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(start_dynamic_module) + assert re.search( + r"for arg\d in range\(arg\d,3,1\)", + utility.dag_builder.get_child_clusters("my_workflow"), + ) + + utility.construct(stop_dynamic_module) + assert re.search( + r"for arg\d in range\(0,arg\d,1\)", + utility.dag_builder.get_child_clusters("my_workflow"), + ) + + utility.construct(step_dynamic_module) + assert re.search( + r"for arg\d in range\(0,3,arg\d\)", + utility.dag_builder.get_child_clusters("my_workflow"), + ) + class TestWhileOp: """Tests that the while loop control flow can be visualized correctly.""" From faf338fc702a2b5b8d4d0fece6d63ce0e02718e5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Mon, 1 Dec 2025 15:47:13 -0500 Subject: [PATCH 189/245] move things around --- .../visualization/test_construct_circuit_dag.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 370e7de4e7..5385fda0f3 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -410,23 +410,23 @@ def start_dynamic(x): qml.H(0) args = (1,) - start_dynamic_module = start_dynamic(*args) - stop_dynamic_module = stop_dynamic(*args) - step_dynamic_module = step_dynamic(*args) - utility = ConstructCircuitDAG(FakeDAGBuilder()) + + start_dynamic_module = start_dynamic(*args) utility.construct(start_dynamic_module) assert re.search( r"for arg\d in range\(arg\d,3,1\)", utility.dag_builder.get_child_clusters("my_workflow"), ) + stop_dynamic_module = stop_dynamic(*args) utility.construct(stop_dynamic_module) assert re.search( r"for arg\d in range\(0,arg\d,1\)", utility.dag_builder.get_child_clusters("my_workflow"), ) + step_dynamic_module = step_dynamic(*args) utility.construct(step_dynamic_module) assert re.search( r"for arg\d in range\(0,3,arg\d\)", From 5d03a81eac04156079548cad7c3be71de0a482df Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 10:21:30 -0500 Subject: [PATCH 190/245] add more qol --- .../python_interface/visualization/construct_circuit_dag.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 2f90832d77..64aef52e7e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -133,7 +133,8 @@ def _if_op(self, operation: scf.IfOp): self.dag_builder.add_cluster( uid, node_label="", - label="", + label="conditional", + labeljust="l", cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) @@ -143,7 +144,7 @@ def _if_op(self, operation: scf.IfOp): uid = f"cluster_ifop_branch{i}_{id(operation)}" self.dag_builder.add_cluster( uid, - node_label=f"if ..." if i == 0 else "else", + node_label="if ..." if i == 0 else "else", label="", cluster_uid=self._cluster_uid_stack[-1], ) From 56e4756898bc19dbe6199e78674ede9c541a4ef7 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 2 Dec 2025 12:51:33 -0500 Subject: [PATCH 191/245] Update frontend/catalyst/python_interface/visualization/construct_circuit_dag.py Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- .../python_interface/visualization/construct_circuit_dag.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 1aeb3ceee7..8291773eb6 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -97,9 +97,6 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: shape="rectangle", ) - for region in operation.regions: - self._visit_region(region) - # ======================= # FuncOp NESTING UTILITY # ======================= From c95c45fcfaea9b2ea1971cc1bc35ca9a7c312cab Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 2 Dec 2025 12:51:40 -0500 Subject: [PATCH 192/245] Update frontend/catalyst/python_interface/visualization/construct_circuit_dag.py Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 8291773eb6..2f54965b13 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -111,7 +111,7 @@ def _func_op(self, operation: func.FuncOp) -> None: label = operation.sym_name.data if "jit_" in operation.sym_name.data: num_qnodes = 0 - for op in operation.walk(): + for op in operation.body.ops: if isinstance(op, catalyst.LaunchKernelOp): num_qnodes += 1 # Get everything after the jit_* prefix From 18aa30b2da36045eca511e14c0b921565437175a Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 2 Dec 2025 12:51:53 -0500 Subject: [PATCH 193/245] Update frontend/catalyst/python_interface/visualization/construct_circuit_dag.py Co-authored-by: Mudit Pandey <18223836+mudit2812@users.noreply.github.com> --- .../python_interface/visualization/construct_circuit_dag.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 2f54965b13..aec969aa17 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -131,8 +131,7 @@ def _func_op(self, operation: func.FuncOp) -> None: ) self._cluster_uid_stack.append(uid) - for region in operation.regions: - self._visit_region(region) + self._visit_block(operation.regions[0].blocks[0]) @_visit_operation.register def _func_return(self, operation: func.ReturnOp) -> None: From 929bd2306eced719dc54927ee7f5614ad7c3386a Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:19:15 -0500 Subject: [PATCH 194/245] add flattened ifop ability --- .../visualization/construct_circuit_dag.py | 43 ++++++++++++++++--- 1 file changed, 37 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 64aef52e7e..e15c9e5592 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,12 +16,11 @@ from functools import singledispatchmethod -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region - from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: @@ -129,6 +128,8 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" + flattened_if_op = _flatten_if_op(operation) + uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( uid, @@ -140,18 +141,30 @@ def _if_op(self, operation: scf.IfOp): self._cluster_uid_stack.append(uid) # Loop through each branch and visualize as a cluster - for i, branch in enumerate(operation.regions): + num_regions = len(flattened_if_op) + for i, (condition_ssa, region) in enumerate(flattened_if_op): + + def _get_conditional_branch_label(i): + if i == 0: + return "if ..." + elif i == num_regions - 1: + return "else" + else: + return "elif ..." + uid = f"cluster_ifop_branch{i}_{id(operation)}" self.dag_builder.add_cluster( uid, - node_label="if ..." if i == 0 else "else", + node_label=_get_conditional_branch_label(i), label="", + style="dashed", + penwidth=1, cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) # Go recursively into the branch to process internals - self._visit_region(branch) + self._visit_region(region) # Pop branch cluster after processing to ensure # logical branches are treated as 'parallel' @@ -228,3 +241,21 @@ def _func_return(self, operation: func.ReturnOp) -> None: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. self._cluster_uid_stack.pop() + + +def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue, Region]]: + """Recursively flattens a nested IfOp (if/elif/else chains).""" + + condition_ssa: SSAValue = op.operands[0] + then_region, else_region = op.regions + + flattened_op: list[tuple[SSAValue, Region]] = [(condition_ssa, then_region)] + + # Peak into else region to see if there's another IfOp + else_block: Block = else_region.block + if isinstance(else_block.ops.last.prev_op, scf.IfOp): + nested_flattened_op = _flatten_if_op(else_block.ops.last.prev_op) + flattened_op.extend(nested_flattened_op) + return flattened_op + + return flattened_op From b51bb7b6cfca74affd41539d9be4180d6287d06b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:20:44 -0500 Subject: [PATCH 195/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 4a79691471..41179fd0f9 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -249,7 +249,10 @@ def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue, Region]]: # Peak into else region to see if there's another IfOp else_block: Block = else_region.block + # Completely relies on the structure that the second last operation + # will be an IfOp (seems to hold true) if isinstance(else_block.ops.last.prev_op, scf.IfOp): + # Recursively flatten any IfOps found in said block nested_flattened_op = _flatten_if_op(else_block.ops.last.prev_op) flattened_op.extend(nested_flattened_op) return flattened_op From 89f45f3bda1d0a39139033126ee5852b88e74b33 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:21:23 -0500 Subject: [PATCH 196/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 41179fd0f9..8704ce4343 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -245,6 +245,7 @@ def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue, Region]]: condition_ssa: SSAValue = op.operands[0] then_region, else_region = op.regions + # Save condition SSA in case we want to visualize it eventually flattened_op: list[tuple[SSAValue, Region]] = [(condition_ssa, then_region)] # Peak into else region to see if there's another IfOp From b45b2f035ce64ed817d4c47df8b7e7d1e6dca592 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:21:48 -0500 Subject: [PATCH 197/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 8704ce4343..5761d2ecc1 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -128,7 +128,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - flattened_if_op = _flatten_if_op(operation) + flattened_if_op: list[tuple[SSAValue, Region]] = _flatten_if_op(operation) uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( From 2028fa743b0d24333fbb1919ee37fd2e5dbcaee8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:22:50 -0500 Subject: [PATCH 198/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 5761d2ecc1..624fab6445 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -170,6 +170,7 @@ def _get_conditional_branch_label(i): # logical branches are treated as 'parallel' self._cluster_uid_stack.pop() + # Pop IfOp cluster before leaving this handler self._cluster_uid_stack.pop() # ============ From 4e994e650ca675a19ac3dab33c61389ab4ffdcf1 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:31:05 -0500 Subject: [PATCH 199/245] simplify PR by making it for ... --- .../visualization/construct_circuit_dag.py | 10 +-- .../test_construct_circuit_dag.py | 75 ++++--------------- 2 files changed, 17 insertions(+), 68 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 624fab6445..e3cb44e118 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -86,18 +86,10 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" - lower_bound, upper_bound, step = ( - resolve_constant_params(operation.lb), - resolve_constant_params(operation.ub), - resolve_constant_params(operation.step), - ) - - index_var_name = operation.body.blocks[0].args[0].name_hint - uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( uid, - node_label=f"for {index_var_name} in range({lower_bound},{upper_bound},{step})", + node_label=f"for ...", label="", cluster_uid=self._cluster_uid_stack[-1], ) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 5385fda0f3..aa4aaf34cc 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,15 +23,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -97,7 +96,9 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_uid"] == parent_cluster_uid: - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] + cluster_label = ( + cluster_data["cluster_label"] or cluster_data["node_label"] + ) cluster_labels.append(cluster_label) return cluster_labels @@ -154,7 +155,9 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph + builder.add_cluster( + "c0", label="Company", cluster_uid=None + ) # Add to base graph builder.add_cluster("c1", label="Marketing", cluster_uid="c0") builder.add_cluster("c2", label="Finance", cluster_uid="c0") @@ -343,7 +346,7 @@ def my_workflow(): utility.construct(module) assert re.search( - r"for arg\d in range\(0,3,1\)", + r"for \.\.\.", utility.dag_builder.get_child_clusters("my_workflow"), ) @@ -370,7 +373,7 @@ def my_workflow(): child_clusters = utility.dag_builder.get_child_clusters("my_workflow") assert len(child_clusters) == 1 assert re.search( - r"for arg\d in range\(0,5,2\)", + r"for \.\.\.", child_clusters, ) @@ -379,60 +382,10 @@ def my_workflow(): child_clusters = utility.dag_builder.get_child_clusters(for_loop_label) assert len(child_clusters) == 1 assert re.search( - r"for arg\d in range\(1,6,2\)", + r"for \.\.\.", child_clusters, ) - def test_dynamic_start_stop_step(self): - """Tests that dynamic start, stop, step variables can be displayed.""" - - dev = qml.device("null.qubit", wires=1) - - @xdsl_from_qjit - @qml.qjit(autograph=True, target="mlir") - @qml.qnode(dev) - def step_dynamic(x): - for i in range(0, 3, x): - qml.H(0) - - @xdsl_from_qjit - @qml.qjit(autograph=True, target="mlir") - @qml.qnode(dev) - def stop_dynamic(x): - for i in range(0, x, 1): - qml.H(0) - - @xdsl_from_qjit - @qml.qjit(autograph=True, target="mlir") - @qml.qnode(dev) - def start_dynamic(x): - for i in range(x, 3, 1): - qml.H(0) - - args = (1,) - utility = ConstructCircuitDAG(FakeDAGBuilder()) - - start_dynamic_module = start_dynamic(*args) - utility.construct(start_dynamic_module) - assert re.search( - r"for arg\d in range\(arg\d,3,1\)", - utility.dag_builder.get_child_clusters("my_workflow"), - ) - - stop_dynamic_module = stop_dynamic(*args) - utility.construct(stop_dynamic_module) - assert re.search( - r"for arg\d in range\(0,arg\d,1\)", - utility.dag_builder.get_child_clusters("my_workflow"), - ) - - step_dynamic_module = step_dynamic(*args) - utility.construct(step_dynamic_module) - assert re.search( - r"for arg\d in range\(0,3,arg\d\)", - utility.dag_builder.get_child_clusters("my_workflow"), - ) - class TestWhileOp: """Tests that the while loop control flow can be visualized correctly.""" @@ -508,6 +461,10 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) + assert "conditional" in utility.dag_builder.get_child_clusters("my_workflow") + assert "if ..." in utility.dag_builder.get_child_clusters("conditional") + assert "elif ..." in utility.dag_builder.get_child_clusters("conditional") + assert "else" in utility.dag_builder.get_child_clusters("conditional") class TestDeviceNode: """Tests that the device node is correctly visualized.""" From e5908470f089230136401a50eeb22d49d2d3ee5a Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 13:32:37 -0500 Subject: [PATCH 200/245] format --- .../visualization/construct_circuit_dag.py | 5 +++-- .../visualization/test_construct_circuit_dag.py | 16 +++++++--------- 2 files changed, 10 insertions(+), 11 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index e3cb44e118..0b5c717a4a 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,11 +16,12 @@ from functools import singledispatchmethod +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue + from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index aa4aaf34cc..29b7b85da3 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,14 +23,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -96,9 +97,7 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_uid"] == parent_cluster_uid: - cluster_label = ( - cluster_data["cluster_label"] or cluster_data["node_label"] - ) + cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] cluster_labels.append(cluster_label) return cluster_labels @@ -155,9 +154,7 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster( - "c0", label="Company", cluster_uid=None - ) # Add to base graph + builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph builder.add_cluster("c1", label="Marketing", cluster_uid="c0") builder.add_cluster("c2", label="Finance", cluster_uid="c0") @@ -466,6 +463,7 @@ def my_workflow(): assert "elif ..." in utility.dag_builder.get_child_clusters("conditional") assert "else" in utility.dag_builder.get_child_clusters("conditional") + class TestDeviceNode: """Tests that the device node is correctly visualized.""" From eb6a0c334311f939deb31e2e8d037381e676967c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:04:42 -0500 Subject: [PATCH 201/245] fix tests --- .../test_construct_circuit_dag.py | 53 ++++++++----------- 1 file changed, 21 insertions(+), 32 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 29b7b85da3..3312078107 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,15 +23,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -97,7 +96,9 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_uid"] == parent_cluster_uid: - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] + cluster_label = ( + cluster_data["cluster_label"] or cluster_data["node_label"] + ) cluster_labels.append(cluster_label) return cluster_labels @@ -154,7 +155,9 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph + builder.add_cluster( + "c0", label="Company", cluster_uid=None + ) # Add to base graph builder.add_cluster("c1", label="Marketing", cluster_uid="c0") builder.add_cluster("c2", label="Finance", cluster_uid="c0") @@ -342,10 +345,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert re.search( - r"for \.\.\.", - utility.dag_builder.get_child_clusters("my_workflow"), - ) + assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") @pytest.mark.unit def test_nested_loop(self): @@ -367,21 +367,10 @@ def my_workflow(): utility.construct(module) # Check first for loop - child_clusters = utility.dag_builder.get_child_clusters("my_workflow") - assert len(child_clusters) == 1 - assert re.search( - r"for \.\.\.", - child_clusters, - ) + assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") # Check second for loop - for_loop_label = utility.dag_builder.get_child_clusters(child_clusters[0]) - child_clusters = utility.dag_builder.get_child_clusters(for_loop_label) - assert len(child_clusters) == 1 - assert re.search( - r"for \.\.\.", - child_clusters, - ) + assert "for ..." in utility.dag_builder.get_child_clusters("for ...") class TestWhileOp: @@ -420,14 +409,14 @@ def test_basic_example(self): @xdsl_from_qjit @qml.qjit(autograph=True, target="mlir") @qml.qnode(dev) - def my_workflow(): - flag = 1 - if flag == 1: + def my_workflow(x): + if x == 2: qml.X(0) else: qml.Y(0) - module = my_workflow() + args = (1,) + module = my_workflow(*args) utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) @@ -444,16 +433,16 @@ def test_if_elif_else_conditional(self): @xdsl_from_qjit @qml.qjit(autograph=True, target="mlir") @qml.qnode(dev) - def my_workflow(): - flag = 1 - if flag == 1: + def my_workflow(x): + if x == 1: qml.X(0) - elif flag == 2: + elif x == 2: qml.Y(0) else: qml.Z(0) - module = my_workflow() + args = (1,) + module = my_workflow(*args) utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) From 7eb31b61dcb07b8401526a94ea6d07081d757a58 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:08:15 -0500 Subject: [PATCH 202/245] format --- .../visualization/test_construct_circuit_dag.py | 15 ++++++--------- 1 file changed, 6 insertions(+), 9 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3312078107..393c55a15e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,14 +23,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -96,9 +97,7 @@ def get_child_clusters(self, parent_cluster_label: str) -> list[str]: cluster_labels = [] for cluster_data in self._clusters.values(): if cluster_data["parent_cluster_uid"] == parent_cluster_uid: - cluster_label = ( - cluster_data["cluster_label"] or cluster_data["node_label"] - ) + cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] cluster_labels.append(cluster_label) return cluster_labels @@ -155,9 +154,7 @@ def builder_with_data(self): builder = FakeDAGBuilder() # Cluster set-up - builder.add_cluster( - "c0", label="Company", cluster_uid=None - ) # Add to base graph + builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph builder.add_cluster("c1", label="Marketing", cluster_uid="c0") builder.add_cluster("c2", label="Finance", cluster_uid="c0") From 8e787c9aa989987477fcb8902a7934f82dd81d6b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:36:03 -0500 Subject: [PATCH 203/245] fix logic in flatten if op --- .../visualization/construct_circuit_dag.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 0b5c717a4a..9a076b9302 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,12 +16,11 @@ from functools import singledispatchmethod -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region, SSAValue - from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: @@ -233,14 +232,14 @@ def _func_return(self, operation: func.ReturnOp) -> None: self._cluster_uid_stack.pop() -def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue, Region]]: +def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue | None, Region]]: """Recursively flattens a nested IfOp (if/elif/else chains).""" condition_ssa: SSAValue = op.operands[0] then_region, else_region = op.regions # Save condition SSA in case we want to visualize it eventually - flattened_op: list[tuple[SSAValue, Region]] = [(condition_ssa, then_region)] + flattened_op: list[tuple[SSAValue | None, Region]] = [(condition_ssa, then_region)] # Peak into else region to see if there's another IfOp else_block: Block = else_region.block @@ -252,4 +251,5 @@ def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue, Region]]: flattened_op.extend(nested_flattened_op) return flattened_op + flattened_op.extend([(None, else_region)]) return flattened_op From 8b6287f7ed632a36b2f7419a023519038e4cfc75 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:36:22 -0500 Subject: [PATCH 204/245] format --- .../python_interface/visualization/construct_circuit_dag.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 9a076b9302..059266e299 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,11 +16,12 @@ from functools import singledispatchmethod +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue + from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: From 83ebb1828014281d94d72c68ce3a8a43972de890 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:37:42 -0500 Subject: [PATCH 205/245] fix type hinting --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 059266e299..1126488417 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -121,7 +121,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - flattened_if_op: list[tuple[SSAValue, Region]] = _flatten_if_op(operation) + flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op(operation) uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( From 6d84d3615c6cf81f12cd94d52da377eb2bec8d6e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 14:38:41 -0500 Subject: [PATCH 206/245] add dev comment --- .../python_interface/visualization/construct_circuit_dag.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 1126488417..2e69eb8b0f 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -252,5 +252,7 @@ def _flatten_if_op(op: scf.IfOp) -> list[tuple[SSAValue | None, Region]]: flattened_op.extend(nested_flattened_op) return flattened_op + # No more nested IfOps, therefore append final region + # with no SSAValue flattened_op.extend([(None, else_region)]) return flattened_op From 3c42841f888ecf5a0cd3d922dbcf38de4b9c58db Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 15:58:55 -0500 Subject: [PATCH 207/245] more deviceop tests --- .../test_construct_circuit_dag.py | 110 +++++++++--------- 1 file changed, 55 insertions(+), 55 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 393c55a15e..9d7f9cd492 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -321,6 +321,61 @@ def my_workflow(): assert "my_qnode2" in utility.dag_builder.get_child_clusters("my_workflow") +class TestDeviceNode: + """Tests that the device node is correctly visualized.""" + + def test_standard_qnode(self): + """Tests that a standard setup works.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + qml.H(0) + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + # Check that device node is within the my_workflow cluster + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_workflow") + assert "NullQubit" in nodes_in_my_workflow + + def test_nested_qnodes(self): + """Tests that nested QJIT'd QNodes are visualized correctly""" + + dev1 = qml.device("null.qubit", wires=1) + dev2 = qml.device("lightning.qubit", wires=1) + + @qml.qnode(dev2) + def my_qnode2(): + qml.X(0) + + @qml.qnode(dev1) + def my_qnode1(): + qml.H(0) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + def my_workflow(): + my_qnode1() + my_qnode2() + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + # Check that device node is within the my_workflow cluster + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode1") + assert "NullQubit" in nodes_in_my_workflow + nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode2") + assert "LightningSimulator" in nodes_in_my_workflow + + class TestForOp: """Tests that the for loop control flow can be visualized correctly.""" @@ -448,58 +503,3 @@ def my_workflow(x): assert "if ..." in utility.dag_builder.get_child_clusters("conditional") assert "elif ..." in utility.dag_builder.get_child_clusters("conditional") assert "else" in utility.dag_builder.get_child_clusters("conditional") - - -class TestDeviceNode: - """Tests that the device node is correctly visualized.""" - - def test_standard_qnode(self): - """Tests that a standard setup works.""" - - dev = qml.device("null.qubit", wires=1) - - @xdsl_from_qjit - @qml.qjit(autograph=True, target="mlir") - @qml.qnode(dev) - def my_workflow(): - qml.H(0) - - module = my_workflow() - - utility = ConstructCircuitDAG(FakeDAGBuilder()) - utility.construct(module) - - # Check that device node is within the my_workflow cluster - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_workflow") - assert "NullQubit" in nodes_in_my_workflow - - def test_nested_qnodes(self): - """Tests that nested QJIT'd QNodes are visualized correctly""" - - dev1 = qml.device("null.qubit", wires=1) - dev2 = qml.device("lightning.qubit", wires=1) - - @qml.qnode(dev2) - def my_qnode2(): - qml.X(0) - - @qml.qnode(dev1) - def my_qnode1(): - qml.H(0) - - @xdsl_from_qjit - @qml.qjit(autograph=True, target="mlir") - def my_workflow(): - my_qnode1() - my_qnode2() - - module = my_workflow() - - utility = ConstructCircuitDAG(FakeDAGBuilder()) - utility.construct(module) - - # Check that device node is within the my_workflow cluster - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode1") - assert "NullQubit" in nodes_in_my_workflow - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode2") - assert "LightningSimulator" in nodes_in_my_workflow From e319aae1b8c19aa5d48fefddffcf2a3eb4aa439e Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Tue, 2 Dec 2025 15:59:18 -0500 Subject: [PATCH 208/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 9d7f9cd492..cf7539b106 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -13,7 +13,6 @@ # limitations under the License. """Unit tests for the ConstructCircuitDAG utility.""" -import re from unittest.mock import Mock import pytest From 1c9c8576661b476f19833bf73f5a1bfc54194540 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 16:04:44 -0500 Subject: [PATCH 209/245] re-order changelog --- doc/releases/changelog-dev.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index 0cb4423675..f9e0be4105 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -6,8 +6,8 @@ [(#2213)](https://github.com/PennyLaneAI/catalyst/pull/2213) [(#2229)](https://github.com/PennyLaneAI/catalyst/pull/2229) [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214) - [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231) [(#2246)](https://github.com/PennyLaneAI/catalyst/pull/2246) + [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From 2ca7bb8fb8b5032f00447b080c8953088357156e Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 16:05:14 -0500 Subject: [PATCH 210/245] add cl entry --- doc/releases/changelog-dev.md | 1 + 1 file changed, 1 insertion(+) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index f9e0be4105..fcdb9f2ed5 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -8,6 +8,7 @@ [(#2214)](https://github.com/PennyLaneAI/catalyst/pull/2214) [(#2246)](https://github.com/PennyLaneAI/catalyst/pull/2246) [(#2231)](https://github.com/PennyLaneAI/catalyst/pull/2231) + [(#2234)](https://github.com/PennyLaneAI/catalyst/pull/2234) * Added ``catalyst.switch``, a qjit compatible, index-switch style control flow decorator. [(#2171)](https://github.com/PennyLaneAI/catalyst/pull/2171) From 61b70d9229a4a2e3c494103c35e6001240f9320c Mon Sep 17 00:00:00 2001 From: andrijapau Date: Tue, 2 Dec 2025 16:14:26 -0500 Subject: [PATCH 211/245] add smaller fontsize for the conditional bit --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 2e69eb8b0f..3faa3cfdcc 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -128,6 +128,7 @@ def _if_op(self, operation: scf.IfOp): uid, node_label="", label="conditional", + fontsize=10, labeljust="l", cluster_uid=self._cluster_uid_stack[-1], ) From 089cb14a4314d5adb261a5c20dbf14a92f8ccd14 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 3 Dec 2025 14:13:42 -0500 Subject: [PATCH 212/245] both single and multi qnode have qjit bounding box --- .../visualization/construct_circuit_dag.py | 33 ++++++------------- .../test_construct_circuit_dag.py | 26 ++++++++------- 2 files changed, 25 insertions(+), 34 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index aec969aa17..e99607b708 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -105,31 +105,18 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: def _func_op(self, operation: func.FuncOp) -> None: """Visit a FuncOp Operation.""" - # If this is the jit_* FuncOp, only draw if there's more than one qnode (launch kernel) - # This avoids redundant nested clusters: jit_my_circuit -> my_circuit -> ... - visualize = True label = operation.sym_name.data if "jit_" in operation.sym_name.data: - num_qnodes = 0 - for op in operation.body.ops: - if isinstance(op, catalyst.LaunchKernelOp): - num_qnodes += 1 - # Get everything after the jit_* prefix - label = str(label).split("_", maxsplit=1)[-1] - if num_qnodes == 1: - visualize = False - - if visualize: - uid = f"cluster_{id(operation)}" - parent_cluster_uid = ( - None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] - ) - self.dag_builder.add_cluster( - uid, - label=label, - cluster_uid=parent_cluster_uid, - ) - self._cluster_uid_stack.append(uid) + label = "qjit" + + uid = f"cluster_{id(operation)}" + parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] + self.dag_builder.add_cluster( + uid, + label=label, + cluster_uid=parent_cluster_uid, + ) + self._cluster_uid_stack.append(uid) self._visit_block(operation.regions[0].blocks[0]) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 051a31e690..7708309c27 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -266,11 +266,15 @@ def my_workflow(): # Check nesting is correct # graph - # └── my_workflow + # └── qjit + # └── my_workflow - # Check my_workflow is nested under my_workflow - my_workflow_id = utility.dag_builder.get_cluster_uid_by_label("my_workflow") - assert graph_clusters[my_workflow_id]["parent_cluster_uid"] == "base" + # Check qjit is nested under graph + qjit_cluster_uid = utility.dag_builder.get_cluster_uid_by_label("qjit") + assert graph_clusters[qjit_cluster_uid]["parent_cluster_uid"] == "base" + + # Check that my_workflow is under qjit + assert "my_workflow" in utility.dag_builder.get_child_clusters("qjit") def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -301,23 +305,23 @@ def my_workflow(): # Check labels we expected are there as clusters graph_clusters = utility.dag_builder.clusters all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - assert "my_workflow" in all_cluster_labels + assert "qjit" in all_cluster_labels assert "my_qnode1" in all_cluster_labels assert "my_qnode2" in all_cluster_labels # Check nesting is correct # graph - # └── my_workflow + # └── qjit # ├── my_qnode1 # └── my_qnode2 - # Check my_workflow is under graph - my_workflow_id = utility.dag_builder.get_cluster_uid_by_label("my_workflow") - assert graph_clusters[my_workflow_id]["parent_cluster_uid"] == "base" + # Check qjit is under graph + qjit_cluster_uid = utility.dag_builder.get_cluster_uid_by_label("qjit") + assert graph_clusters[qjit_cluster_uid]["parent_cluster_uid"] == "base" # Check both qnodes are under my_workflow - assert "my_qnode1" in utility.dag_builder.get_child_clusters("my_workflow") - assert "my_qnode2" in utility.dag_builder.get_child_clusters("my_workflow") + assert "my_qnode1" in utility.dag_builder.get_child_clusters("qjit") + assert "my_qnode2" in utility.dag_builder.get_child_clusters("qjit") class TestDeviceNode: From 2a2d55c39afe675c0107de3ce02bc9756fad39ba Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Wed, 3 Dec 2025 14:14:50 -0500 Subject: [PATCH 213/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 7708309c27..b656bd34ee 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -262,6 +262,7 @@ def my_workflow(): # Check labels we expected are there graph_clusters = utility.dag_builder.clusters all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} + assert "qjit" in all_cluster_labels assert "my_workflow" in all_cluster_labels # Check nesting is correct From f7261a982b478178cc6e04ee5e0ed01f6143fd1d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 3 Dec 2025 14:23:21 -0500 Subject: [PATCH 214/245] add more detail to the for loop label --- .../visualization/construct_circuit_dag.py | 20 +++++++++++++------ 1 file changed, 14 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index da94790bb0..ea4ebf8036 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,12 +16,11 @@ from functools import singledispatchmethod -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region, SSAValue - from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: @@ -88,9 +87,14 @@ def _visit_block(self, block: Block) -> None: def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" uid = f"cluster_{id(operation)}" + + # TODO: Extract from IR in future PR + iter_var = "..." + start, stop, step = "...", "...", "..." + label = f"for {iter_var} in range({start}, {stop}, {step})" self.dag_builder.add_cluster( uid, - node_label=f"for ...", + node_label=label, label="", cluster_uid=self._cluster_uid_stack[-1], ) @@ -121,7 +125,9 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op(operation) + flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op( + operation + ) uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( @@ -198,7 +204,9 @@ def _func_op(self, operation: func.FuncOp) -> None: label = "qjit" uid = f"cluster_{id(operation)}" - parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] + parent_cluster_uid = ( + None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] + ) self.dag_builder.add_cluster( uid, label=label, From a8cd49c32765afef31e202c3e33c27cca18a0d37 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 3 Dec 2025 14:26:03 -0500 Subject: [PATCH 215/245] format --- .../visualization/construct_circuit_dag.py | 13 +++++-------- 1 file changed, 5 insertions(+), 8 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index ea4ebf8036..9e628958b1 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,11 +16,12 @@ from functools import singledispatchmethod +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue + from catalyst.python_interface.dialects import catalyst, quantum from catalyst.python_interface.inspection.xdsl_conversion import resolve_constant_params from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import builtin, func, scf -from xdsl.ir import Block, Operation, Region, SSAValue class ConstructCircuitDAG: @@ -125,9 +126,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: @_visit_operation.register def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" - flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op( - operation - ) + flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op(operation) uid = f"cluster_{id(operation)}" self.dag_builder.add_cluster( @@ -204,9 +203,7 @@ def _func_op(self, operation: func.FuncOp) -> None: label = "qjit" uid = f"cluster_{id(operation)}" - parent_cluster_uid = ( - None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] - ) + parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] self.dag_builder.add_cluster( uid, label=label, From 70b0d937ff157bce8a7424c93606be1047b1b0b4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Wed, 3 Dec 2025 14:49:35 -0500 Subject: [PATCH 216/245] fix test --- .../visualization/test_construct_circuit_dag.py | 12 +++++++++--- 1 file changed, 9 insertions(+), 3 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index eeeb6a996f..ac01bb020c 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -401,7 +401,9 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") + assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( + "my_workflow" + ) @pytest.mark.unit def test_nested_loop(self): @@ -423,10 +425,14 @@ def my_workflow(): utility.construct(module) # Check first for loop - assert "for ..." in utility.dag_builder.get_child_clusters("my_workflow") + assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( + "my_workflow" + ) # Check second for loop - assert "for ..." in utility.dag_builder.get_child_clusters("for ...") + assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( + "for ... in range(..., ..., ...)" + ) class TestWhileOp: From cb0592190624c3771e112c8af539d77bb4ff689d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:11:23 -0500 Subject: [PATCH 217/245] use counter instead of builtin id function --- .../visualization/construct_circuit_dag.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index e99607b708..ec278ef34e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -44,6 +44,9 @@ def __init__(self, dag_builder: DAGBuilder) -> None: # Keep track of nesting clusters using a stack self._cluster_uid_stack: list[str] = [] + # Use counter internally for UID + self._uid_counter = 0 + def _reset(self) -> None: """Resets the instance.""" self._cluster_uid_stack: list[str] = [] @@ -86,7 +89,8 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _device_init(self, operation: quantum.DeviceInitOp) -> None: """Handles the initialization of a quantum device.""" - node_id = f"node_{id(operation)}" + node_id = f"device_node{self._uid_counter}" + self._uid_counter += 1 self.dag_builder.add_node( node_id, label=operation.device_name.data, @@ -109,7 +113,8 @@ def _func_op(self, operation: func.FuncOp) -> None: if "jit_" in operation.sym_name.data: label = "qjit" - uid = f"cluster_{id(operation)}" + uid = f"funcop_cluster_{self._uid_counter}" + self._uid_counter += 1 parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] self.dag_builder.add_cluster( uid, From 1883fb3e0a4a55dc4f8369d3819d034e111e7ca5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:12:02 -0500 Subject: [PATCH 218/245] whoops --- .../python_interface/visualization/construct_circuit_dag.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index ec278ef34e..43a0089bff 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -45,11 +45,12 @@ def __init__(self, dag_builder: DAGBuilder) -> None: self._cluster_uid_stack: list[str] = [] # Use counter internally for UID - self._uid_counter = 0 + self._uid_counter: int = 0 def _reset(self) -> None: """Resets the instance.""" self._cluster_uid_stack: list[str] = [] + self._uid_counter: int = 0 def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. From 5e2d456607d93392ecbe4795e667524d7421f647 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 4 Dec 2025 11:13:16 -0500 Subject: [PATCH 219/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 43a0089bff..c478aadfc4 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -90,7 +90,7 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _device_init(self, operation: quantum.DeviceInitOp) -> None: """Handles the initialization of a quantum device.""" - node_id = f"device_node{self._uid_counter}" + node_id = f"device_node_{self._uid_counter}" self._uid_counter += 1 self.dag_builder.add_node( node_id, From 36b5f7d15d261edee392ab88c712600336610068 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 4 Dec 2025 11:17:21 -0500 Subject: [PATCH 220/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index c478aadfc4..b48c071517 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -131,8 +131,6 @@ def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first cluster as it is the "base" of the graph diagram. - # If it is a multi-qnode workflow, it will represent the "workflow" function - # If it is a single qnode, it will represent the quantum function. if len(self._cluster_uid_stack) > 1: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. From ee15b67c7b43e97590afdcd92a20b7bd1fb82ab3 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 4 Dec 2025 11:18:05 -0500 Subject: [PATCH 221/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/construct_circuit_dag.py | 1 + 1 file changed, 1 insertion(+) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index b48c071517..85dee3eb1c 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -131,6 +131,7 @@ def _func_return(self, operation: func.ReturnOp) -> None: """Handle func.return to exit FuncOp's cluster scope.""" # NOTE: Skip first cluster as it is the "base" of the graph diagram. + # In our case, it is the `qjit` bounding box. if len(self._cluster_uid_stack) > 1: # If we hit a func.return operation we know we are leaving # the FuncOp's scope and so we can pop the ID off the stack. From 16e003d39998f26ea2ad8e04558b91f88fa8f31b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:22:24 -0500 Subject: [PATCH 222/245] add two counters --- .../visualization/construct_circuit_dag.py | 14 ++++++++------ 1 file changed, 8 insertions(+), 6 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 85dee3eb1c..ef13948b18 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -45,12 +45,14 @@ def __init__(self, dag_builder: DAGBuilder) -> None: self._cluster_uid_stack: list[str] = [] # Use counter internally for UID - self._uid_counter: int = 0 + self._node_uid_counter: int = 0 + self._cluster_uid_counter: int = 0 def _reset(self) -> None: """Resets the instance.""" self._cluster_uid_stack: list[str] = [] - self._uid_counter: int = 0 + self._node_uid_counter: int = 0 + self._cluster_uid_counter: int = 0 def construct(self, module: builtin.ModuleOp) -> None: """Constructs the DAG from the module. @@ -90,8 +92,7 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _device_init(self, operation: quantum.DeviceInitOp) -> None: """Handles the initialization of a quantum device.""" - node_id = f"device_node_{self._uid_counter}" - self._uid_counter += 1 + node_id = f"node{self._node_uid_counter}" self.dag_builder.add_node( node_id, label=operation.device_name.data, @@ -101,6 +102,7 @@ def _device_init(self, operation: quantum.DeviceInitOp) -> None: penwidth=2, shape="rectangle", ) + self._node_uid_counter += 1 # ======================= # FuncOp NESTING UTILITY @@ -114,14 +116,14 @@ def _func_op(self, operation: func.FuncOp) -> None: if "jit_" in operation.sym_name.data: label = "qjit" - uid = f"funcop_cluster_{self._uid_counter}" - self._uid_counter += 1 + uid = f"cluster{self._cluster_uid_counter}" parent_cluster_uid = None if self._cluster_uid_stack == [] else self._cluster_uid_stack[-1] self.dag_builder.add_cluster( uid, label=label, cluster_uid=parent_cluster_uid, ) + self._cluster_uid_counter += 1 self._cluster_uid_stack.append(uid) self._visit_block(operation.regions[0].blocks[0]) From d5ccaf7f91876e464d60dd5a40c0e927a64d56c3 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:40:37 -0500 Subject: [PATCH 223/245] use the counters for uid --- .../visualization/construct_circuit_dag.py | 13 +++++++++---- 1 file changed, 9 insertions(+), 4 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 51870da586..4f63fe26a4 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -93,12 +93,13 @@ def _visit_block(self, block: Block) -> None: @_visit_operation.register def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" - uid = f"cluster_{id(operation)}" # TODO: Extract from IR in future PR iter_var = "..." start, stop, step = "...", "...", "..." label = f"for {iter_var} in range({start}, {stop}, {step})" + + uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, node_label=label, @@ -106,6 +107,7 @@ def _for_op(self, operation: scf.ForOp) -> None: cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) + self._cluster_uid_counter += 1 for region in operation.regions: self._visit_region(region) @@ -115,7 +117,7 @@ def _for_op(self, operation: scf.ForOp) -> None: @_visit_operation.register def _while_op(self, operation: scf.WhileOp) -> None: """Handle an xDSL WhileOp operation.""" - uid = f"cluster_{id(operation)}" + uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, node_label="while ...", @@ -123,6 +125,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) + self._cluster_uid_counter += 1 for region in operation.regions: self._visit_region(region) @@ -134,7 +137,7 @@ def _if_op(self, operation: scf.IfOp): """Handles the scf.IfOp operation.""" flattened_if_op: list[tuple[SSAValue | None, Region]] = _flatten_if_op(operation) - uid = f"cluster_{id(operation)}" + uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, node_label="", @@ -144,6 +147,7 @@ def _if_op(self, operation: scf.IfOp): cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) + self._cluster_uid_counter += 1 # Loop through each branch and visualize as a cluster num_regions = len(flattened_if_op) @@ -157,7 +161,7 @@ def _get_conditional_branch_label(i): else: return "elif ..." - uid = f"cluster_ifop_branch{i}_{id(operation)}" + uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, node_label=_get_conditional_branch_label(i), @@ -167,6 +171,7 @@ def _get_conditional_branch_label(i): cluster_uid=self._cluster_uid_stack[-1], ) self._cluster_uid_stack.append(uid) + self._cluster_uid_counter += 1 # Go recursively into the branch to process internals self._visit_region(region) From 1dd9f38797651010857d70fc4f43f044b7e23cdd Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:56:52 -0500 Subject: [PATCH 224/245] update tests to be simpler --- .../test_construct_circuit_dag.py | 185 +++++------------- 1 file changed, 48 insertions(+), 137 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index b656bd34ee..06e5d99912 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -77,53 +77,6 @@ def add_cluster( "attrs": attrs, } - def get_nodes_in_cluster(self, cluster_label: str) -> list[str]: - """ - Returns a list of node labels that are direct children of the given cluster. - """ - node_uids = [] - cluster_uid = self.get_cluster_uid_by_label(cluster_label) - for node_data in self._nodes.values(): - if node_data["parent_cluster_uid"] == cluster_uid: - node_uids.append(node_data["label"]) - return node_uids - - def get_child_clusters(self, parent_cluster_label: str) -> list[str]: - """ - Returns a list of cluster labels that are direct children of the given parent cluster. - """ - parent_cluster_uid = self.get_cluster_uid_by_label(parent_cluster_label) - cluster_labels = [] - for cluster_data in self._clusters.values(): - if cluster_data["parent_cluster_uid"] == parent_cluster_uid: - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] - cluster_labels.append(cluster_label) - return cluster_labels - - def get_node_uid_by_label(self, label: str) -> str | None: - """ - Finds the ID of a node given its label. - Assumes labels are unique for testing purposes. - """ - for id, node_data in self._nodes.items(): - if node_data["label"] == label: - return id - return None - - def get_cluster_uid_by_label(self, label: str) -> str | None: - """ - Finds the ID of a cluster given its label. - Assumes cluster labels are unique for testing purposes. - """ - # Work around for base graph - if label == "base": - return "base" - for id, cluster_data in self._clusters.items(): - cluster_label = cluster_data["cluster_label"] or cluster_data["node_label"] - if cluster_label == label: - return id - return None - @property def nodes(self): return self._nodes @@ -143,71 +96,6 @@ def to_string(self) -> str: return "graph" -class TestFakeDAGBuilder: - """Test the FakeDAGBuilder to ensure helper functions work as intended.""" - - @pytest.fixture - def builder_with_data(self): - """Sets up an instance with a complex graph already built.""" - - builder = FakeDAGBuilder() - - # Cluster set-up - builder.add_cluster("c0", label="Company", cluster_uid=None) # Add to base graph - builder.add_cluster("c1", label="Marketing", cluster_uid="c0") - builder.add_cluster("c2", label="Finance", cluster_uid="c0") - - # Node set-up - builder.add_node("n0", "CEO", cluster_uid="c0") - builder.add_node("n1", "Marketing Manager", cluster_uid="c1") - builder.add_node("n2", "Finance Manager", cluster_uid="c2") - - return builder - - # Test ID look up - - def test_get_node_uid_by_label_success(self, builder_with_data): - assert builder_with_data.get_node_uid_by_label("Finance Manager") == "n2" - assert builder_with_data.get_node_uid_by_label("Marketing Manager") == "n1" - assert builder_with_data.get_node_uid_by_label("CEO") == "n0" - - def test_get_node_uid_by_label_failure(self, builder_with_data): - assert builder_with_data.get_node_uid_by_label("Software Manager") is None - - def test_get_cluster_uid_by_label_success(self, builder_with_data): - assert builder_with_data.get_cluster_uid_by_label("Finance") == "c2" - assert builder_with_data.get_cluster_uid_by_label("Marketing") == "c1" - assert builder_with_data.get_cluster_uid_by_label("Company") == "c0" - - def test_get_cluster_uid_by_label_failure(self, builder_with_data): - assert builder_with_data.get_cluster_uid_by_label("Software") is None - - # Test relationship probing - - def test_node_heirarchy(self, builder_with_data): - finance_nodes = builder_with_data.get_nodes_in_cluster("Finance") - assert finance_nodes == ["Finance Manager"] - - marketing_nodes = builder_with_data.get_nodes_in_cluster("Marketing") - assert marketing_nodes == ["Marketing Manager"] - - company_nodes = builder_with_data.get_nodes_in_cluster("Company") - assert company_nodes == ["CEO"] - - def test_cluster_heirarchy(self, builder_with_data): - clusters_in_finance = builder_with_data.get_child_clusters("Finance") - assert not clusters_in_finance - - clusters_in_marketing = builder_with_data.get_child_clusters("Marketing") - assert not clusters_in_marketing - - clusters_in_company = builder_with_data.get_child_clusters("Company") - assert {"Finance", "Marketing"} == set(clusters_in_company) - - clusters_in_base = builder_with_data.get_child_clusters("base") - assert clusters_in_base == ["Company"] - - @pytest.mark.unit def test_dependency_injection(): """Tests that relevant dependencies are injected.""" @@ -259,11 +147,7 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - # Check labels we expected are there graph_clusters = utility.dag_builder.clusters - all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - assert "qjit" in all_cluster_labels - assert "my_workflow" in all_cluster_labels # Check nesting is correct # graph @@ -271,11 +155,13 @@ def my_workflow(): # └── my_workflow # Check qjit is nested under graph - qjit_cluster_uid = utility.dag_builder.get_cluster_uid_by_label("qjit") - assert graph_clusters[qjit_cluster_uid]["parent_cluster_uid"] == "base" + assert graph_clusters["cluster0"]["cluster_label"] == "qjit" + assert graph_clusters["cluster0"]["parent_cluster_uid"] == "base" # Check that my_workflow is under qjit - assert "my_workflow" in utility.dag_builder.get_child_clusters("qjit") + assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow" + assert graph_clusters["cluster1"]["parent_cluster_uid"] == "cluster0" + def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -303,12 +189,6 @@ def my_workflow(): graph_clusters = utility.dag_builder.clusters - # Check labels we expected are there as clusters - graph_clusters = utility.dag_builder.clusters - all_cluster_labels = {info["cluster_label"] for info in graph_clusters.values()} - assert "qjit" in all_cluster_labels - assert "my_qnode1" in all_cluster_labels - assert "my_qnode2" in all_cluster_labels # Check nesting is correct # graph @@ -317,12 +197,15 @@ def my_workflow(): # └── my_qnode2 # Check qjit is under graph - qjit_cluster_uid = utility.dag_builder.get_cluster_uid_by_label("qjit") - assert graph_clusters[qjit_cluster_uid]["parent_cluster_uid"] == "base" + assert graph_clusters["cluster0"]["cluster_label"] == "qjit" + assert graph_clusters["cluster0"]["parent_cluster_uid"] == "base" # Check both qnodes are under my_workflow - assert "my_qnode1" in utility.dag_builder.get_child_clusters("qjit") - assert "my_qnode2" in utility.dag_builder.get_child_clusters("qjit") + assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1" + assert graph_clusters["cluster1"]["parent_cluster_uid"] == "cluster0" + + assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2" + assert graph_clusters["cluster2"]["parent_cluster_uid"] == "cluster0" class TestDeviceNode: @@ -344,9 +227,20 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - # Check that device node is within the my_workflow cluster - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_workflow") - assert "NullQubit" in nodes_in_my_workflow + graph_nodes = utility.dag_builder.nodes + graph_clusters = utility.dag_builder.clusters + + # Check nesting is correct + # graph + # └── qjit + # └── my_workflow: NullQubit + + # Assert device node is inside my_workflow cluster + assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow" + assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1" + + # Assert label is as expected + assert graph_nodes["node0"]["label"] == "NullQubit" def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -373,8 +267,25 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - # Check that device node is within the my_workflow cluster - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode1") - assert "NullQubit" in nodes_in_my_workflow - nodes_in_my_workflow = utility.dag_builder.get_nodes_in_cluster("my_qnode2") - assert "LightningSimulator" in nodes_in_my_workflow + graph_nodes = utility.dag_builder.nodes + graph_clusters = utility.dag_builder.clusters + + # Check nesting is correct + # graph + # └── qjit + # ├── my_qnode1: NullQubit + # └── my_qnode2: LightningSimulator + + # Assert device node is inside my_workflow cluster + assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1" + assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1" + + # Assert label is as expected + assert graph_nodes["node0"]["label"] == "LightningSimulator" + + # Assert device node is inside my_workflow cluster + assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2" + assert graph_nodes["node1"]["parent_cluster_uid"] == "cluster2" + + # Assert label is as expected + assert graph_nodes["node1"]["label"] == "NullQubit" From 6a9b01dcf85c4f33d2068472a878043faeb11d88 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 11:57:56 -0500 Subject: [PATCH 225/245] fix tests --- .../visualization/test_construct_circuit_dag.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 06e5d99912..858689fd73 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -49,7 +49,7 @@ def add_node(self, uid, label, cluster_uid=None, **attrs) -> None: self._nodes[uid] = { "uid": uid, "label": label, - "parent_cluster_uid": "base" if cluster_uid is None else cluster_uid, + "parent_cluster_uid": cluster_uid, "attrs": attrs, } @@ -73,7 +73,7 @@ def add_cluster( "uid": uid, "node_label": node_label, "cluster_label": attrs.get("label"), - "parent_cluster_uid": "base" if cluster_uid is None else cluster_uid, + "parent_cluster_uid": cluster_uid, "attrs": attrs, } @@ -156,7 +156,7 @@ def my_workflow(): # Check qjit is nested under graph assert graph_clusters["cluster0"]["cluster_label"] == "qjit" - assert graph_clusters["cluster0"]["parent_cluster_uid"] == "base" + assert graph_clusters["cluster0"]["parent_cluster_uid"] is None # Check that my_workflow is under qjit assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow" @@ -198,7 +198,7 @@ def my_workflow(): # Check qjit is under graph assert graph_clusters["cluster0"]["cluster_label"] == "qjit" - assert graph_clusters["cluster0"]["parent_cluster_uid"] == "base" + assert graph_clusters["cluster0"]["parent_cluster_uid"] is None # Check both qnodes are under my_workflow assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1" From a2ac9ff74185b25eeab6c4d6acc929e96e30f0c7 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:00:30 -0500 Subject: [PATCH 226/245] minor fix --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 858689fd73..0f094fe375 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -276,14 +276,14 @@ def my_workflow(): # ├── my_qnode1: NullQubit # └── my_qnode2: LightningSimulator - # Assert device node is inside my_workflow cluster + # Assert lightning.qubit device node is inside my_qnode1 cluster assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1" assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1" # Assert label is as expected assert graph_nodes["node0"]["label"] == "LightningSimulator" - # Assert device node is inside my_workflow cluster + # Assert null qubit device node is inside my_qnode2 cluster assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2" assert graph_nodes["node1"]["parent_cluster_uid"] == "cluster2" From 0dbe4fe8573cadf782ad5a39fd62f84293c40a0d Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:06:11 -0500 Subject: [PATCH 227/245] foramt --- .../visualization/test_construct_circuit_dag.py | 8 +++----- 1 file changed, 3 insertions(+), 5 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 0f094fe375..2c7e637b57 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -49,7 +49,7 @@ def add_node(self, uid, label, cluster_uid=None, **attrs) -> None: self._nodes[uid] = { "uid": uid, "label": label, - "parent_cluster_uid": cluster_uid, + "parent_cluster_uid": cluster_uid, "attrs": attrs, } @@ -156,13 +156,12 @@ def my_workflow(): # Check qjit is nested under graph assert graph_clusters["cluster0"]["cluster_label"] == "qjit" - assert graph_clusters["cluster0"]["parent_cluster_uid"] is None + assert graph_clusters["cluster0"]["parent_cluster_uid"] is None # Check that my_workflow is under qjit assert graph_clusters["cluster1"]["cluster_label"] == "my_workflow" assert graph_clusters["cluster1"]["parent_cluster_uid"] == "cluster0" - def test_nested_qnodes(self): """Tests that nested QJIT'd QNodes are visualized correctly""" @@ -189,7 +188,6 @@ def my_workflow(): graph_clusters = utility.dag_builder.clusters - # Check nesting is correct # graph # └── qjit @@ -198,7 +196,7 @@ def my_workflow(): # Check qjit is under graph assert graph_clusters["cluster0"]["cluster_label"] == "qjit" - assert graph_clusters["cluster0"]["parent_cluster_uid"] is None + assert graph_clusters["cluster0"]["parent_cluster_uid"] is None # Check both qnodes are under my_workflow assert graph_clusters["cluster1"]["cluster_label"] == "my_qnode1" From 8afc67a26782a68852244d43418e98839c592c69 Mon Sep 17 00:00:00 2001 From: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com> Date: Thu, 4 Dec 2025 12:08:23 -0500 Subject: [PATCH 228/245] Apply suggestion from @andrijapau --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index cf27992c57..3759adcc2e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -264,7 +264,6 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - graph_nodes = utility.dag_builder.nodes graph_clusters = utility.dag_builder.clusters From c0d4d67be0ea9030544eb223199c794ea5f2afff Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:16:19 -0500 Subject: [PATCH 229/245] update tests --- .../test_construct_circuit_dag.py | 65 +++++++++++++------ 1 file changed, 46 insertions(+), 19 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3759adcc2e..186a780c97 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -15,6 +15,7 @@ from unittest.mock import Mock +from jax import util import pytest pytestmark = pytest.mark.usefixtures("requires_xdsl") @@ -309,9 +310,11 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( - "my_workflow" - ) + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" @pytest.mark.unit def test_nested_loop(self): @@ -332,15 +335,14 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - # Check first for loop - assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( - "my_workflow" - ) + clusters = utility.dag_builder.clusters - # Check second for loop - assert "for ... in range(..., ..., ...)" in utility.dag_builder.get_child_clusters( - "for ... in range(..., ..., ...)" - ) + # cluster0 -> qjit + # cluster1 -> my_workflow + assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" + assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + assert clusters['cluster3']['node_label'] == "for ... in range(..., ..., ...)" + assert clusters['cluster3']['node_label'] == "cluster2" class TestWhileOp: @@ -365,7 +367,11 @@ def my_workflow(): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "while ..." in utility.dag_builder.get_child_clusters("my_workflow") + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + assert clusters['cluster2']['node_label'] == "while ..." class TestIfOp: @@ -391,9 +397,19 @@ def my_workflow(x): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "conditional" in utility.dag_builder.get_child_clusters("my_workflow") - assert "if ..." in utility.dag_builder.get_child_clusters("conditional") - assert "else" in utility.dag_builder.get_child_clusters("conditional") + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + # Check conditional is a cluster within cluster1 (my_workflow) + assert clusters['cluster2']['cluster_label'] == "conditional" + assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + + # Check three clusters live within cluster2 (conditional) + assert clusters['cluster3']['node_label'] == "if ..." + assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" + assert clusters['cluster5']['node_label'] == "else" + assert clusters['cluster5']['parent_cluster_uid'] == "cluster2" @pytest.mark.unit def test_if_elif_else_conditional(self): @@ -417,7 +433,18 @@ def my_workflow(x): utility = ConstructCircuitDAG(FakeDAGBuilder()) utility.construct(module) - assert "conditional" in utility.dag_builder.get_child_clusters("my_workflow") - assert "if ..." in utility.dag_builder.get_child_clusters("conditional") - assert "elif ..." in utility.dag_builder.get_child_clusters("conditional") - assert "else" in utility.dag_builder.get_child_clusters("conditional") + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + # Check conditional is a cluster within my_workflow + assert clusters['cluster2']['cluster_label'] == "conditional" + assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + + # Check three clusters live within conditional + assert clusters['cluster3']['node_label'] == "if ..." + assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" + assert clusters['cluster4']['node_label'] == "elif ..." + assert clusters['cluster4']['parent_cluster_uid'] == "cluster2" + assert clusters['cluster5']['node_label'] == "else" + assert clusters['cluster5']['parent_cluster_uid'] == "cluster2" From 97069677c85c27ec5ffb2401857757978dcbce70 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:17:15 -0500 Subject: [PATCH 230/245] update tests --- .../visualization/test_construct_circuit_dag.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 186a780c97..f9d5108379 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -315,6 +315,7 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" + assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" @pytest.mark.unit def test_nested_loop(self): @@ -342,7 +343,7 @@ def my_workflow(): assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" assert clusters['cluster3']['node_label'] == "for ... in range(..., ..., ...)" - assert clusters['cluster3']['node_label'] == "cluster2" + assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" class TestWhileOp: @@ -372,6 +373,7 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow assert clusters['cluster2']['node_label'] == "while ..." + assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" class TestIfOp: From 5e62c31ead59540358740c81201afabe11de50bc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:19:00 -0500 Subject: [PATCH 231/245] format --- .../test_construct_circuit_dag.py | 46 +++++++++---------- 1 file changed, 23 insertions(+), 23 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index f9d5108379..d883c5a9fc 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -15,8 +15,8 @@ from unittest.mock import Mock -from jax import util import pytest +from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl") @@ -314,8 +314,8 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" - assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + assert clusters["cluster2"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" @pytest.mark.unit def test_nested_loop(self): @@ -340,10 +340,10 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters['cluster2']['node_label'] == "for ... in range(..., ..., ...)" - assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" - assert clusters['cluster3']['node_label'] == "for ... in range(..., ..., ...)" - assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" + assert clusters["cluster2"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" + assert clusters["cluster3"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" class TestWhileOp: @@ -372,8 +372,8 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters['cluster2']['node_label'] == "while ..." - assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + assert clusters["cluster2"]["node_label"] == "while ..." + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" class TestIfOp: @@ -404,14 +404,14 @@ def my_workflow(x): # cluster0 -> qjit # cluster1 -> my_workflow # Check conditional is a cluster within cluster1 (my_workflow) - assert clusters['cluster2']['cluster_label'] == "conditional" - assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + assert clusters["cluster2"]["cluster_label"] == "conditional" + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" # Check three clusters live within cluster2 (conditional) - assert clusters['cluster3']['node_label'] == "if ..." - assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" - assert clusters['cluster5']['node_label'] == "else" - assert clusters['cluster5']['parent_cluster_uid'] == "cluster2" + assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" + assert clusters["cluster5"]["node_label"] == "else" + assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" @pytest.mark.unit def test_if_elif_else_conditional(self): @@ -440,13 +440,13 @@ def my_workflow(x): # cluster0 -> qjit # cluster1 -> my_workflow # Check conditional is a cluster within my_workflow - assert clusters['cluster2']['cluster_label'] == "conditional" - assert clusters['cluster2']['parent_cluster_uid'] == "cluster1" + assert clusters["cluster2"]["cluster_label"] == "conditional" + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" # Check three clusters live within conditional - assert clusters['cluster3']['node_label'] == "if ..." - assert clusters['cluster3']['parent_cluster_uid'] == "cluster2" - assert clusters['cluster4']['node_label'] == "elif ..." - assert clusters['cluster4']['parent_cluster_uid'] == "cluster2" - assert clusters['cluster5']['node_label'] == "else" - assert clusters['cluster5']['parent_cluster_uid'] == "cluster2" + assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" + assert clusters["cluster4"]["node_label"] == "elif ..." + assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" + assert clusters["cluster5"]["node_label"] == "else" + assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" From 1474a66b9c82c610609344d66c06f0dd881539f5 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:22:27 -0500 Subject: [PATCH 232/245] whoops fix test --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 2c7e637b57..d4cac366a2 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -279,11 +279,11 @@ def my_workflow(): assert graph_nodes["node0"]["parent_cluster_uid"] == "cluster1" # Assert label is as expected - assert graph_nodes["node0"]["label"] == "LightningSimulator" + assert graph_nodes["node0"]["label"] == "NullQubit" # Assert null qubit device node is inside my_qnode2 cluster assert graph_clusters["cluster2"]["cluster_label"] == "my_qnode2" assert graph_nodes["node1"]["parent_cluster_uid"] == "cluster2" # Assert label is as expected - assert graph_nodes["node1"]["label"] == "NullQubit" + assert graph_nodes["node1"]["label"] == "LightningSimulator" From ced6f7550263ebf498bdff76ad3b58d6398e17a8 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:25:16 -0500 Subject: [PATCH 233/245] format --- .../visualization/test_construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index e71e94ec1a..eac990e092 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -449,4 +449,4 @@ def my_workflow(x): assert clusters["cluster4"]["node_label"] == "elif ..." assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" assert clusters["cluster5"]["node_label"] == "else" - assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" \ No newline at end of file + assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" From 5641e680c9514fcad3a7866afaac8b6d6b2f9e61 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:42:55 -0500 Subject: [PATCH 234/245] fix test --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index eac990e092..d563224e5e 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -410,8 +410,8 @@ def my_workflow(x): # Check three clusters live within cluster2 (conditional) assert clusters["cluster3"]["node_label"] == "if ..." assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" - assert clusters["cluster5"]["node_label"] == "else" - assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" + assert clusters["cluster4"]["node_label"] == "else" + assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" @pytest.mark.unit def test_if_elif_else_conditional(self): From 22c8ae190936f99266d6c34d5e889c4563ec894b Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:58:01 -0500 Subject: [PATCH 235/245] add more nested testing --- .../test_construct_circuit_dag.py | 87 ++++++++++++++++++- 1 file changed, 83 insertions(+), 4 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index d563224e5e..787ac2ca74 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,15 +23,14 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region - from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -375,6 +374,38 @@ def my_workflow(): assert clusters["cluster2"]["node_label"] == "while ..." assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" + @pytest.mark.unit + def test_nested_loop(self): + """Tests that nested while loops are visualized correctly.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(): + outer_counter = 0 + inner_counter = 0 + while outer_counter < 5: + while inner_counter < 6: + qml.H(0) + inner_counter += 1 + outer_counter += 1 + + module = my_workflow() + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + assert clusters["cluster2"]["node_label"] == "while ..." + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" + assert clusters["cluster3"]["node_label"] == "while ..." + assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" + class TestIfOp: """Tests that the conditional control flow can be visualized correctly.""" @@ -450,3 +481,51 @@ def my_workflow(x): assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" assert clusters["cluster5"]["node_label"] == "else" assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" + + @pytest.mark.unit + def test_nested_conditionals(self): + """Tests that nested conditionals are visualized correctly.""" + + dev = qml.device("null.qubit", wires=1) + + @xdsl_from_qjit + @qml.qjit(autograph=True, target="mlir") + @qml.qnode(dev) + def my_workflow(x,y): + if x == 1: + if y == 2: + qml.H(0) + else: + qml.Z(0) + qml.X(0) + else: + qml.Z(0) + + args = (1,2) + module = my_workflow(*args) + + utility = ConstructCircuitDAG(FakeDAGBuilder()) + utility.construct(module) + + clusters = utility.dag_builder.clusters + clusters = utility.dag_builder.clusters + + # cluster0 -> qjit + # cluster1 -> my_workflow + + # Check first conditional is a cluster within my_workflow + assert clusters["cluster2"]["cluster_label"] == "conditional" + assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" + + # Check two clusters live within first conditional + assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" + # first conditional's else + assert clusters["cluster6"]["node_label"] == "else" + assert clusters["cluster6"]["parent_cluster_uid"] == "cluster2" + + # Check nested if / else is within the first if cluster + assert clusters["cluster4"]["node_label"] == "if ..." + assert clusters["cluster4"]["parent_cluster_uid"] == "cluster3" + assert clusters["cluster5"]["node_label"] == "if ..." + assert clusters["cluster5"]["parent_cluster_uid"] == "cluster3" From ab7c4714f0ce52e2c219892df796aa726676bfde Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 12:59:11 -0500 Subject: [PATCH 236/245] format --- .../visualization/test_construct_circuit_dag.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 787ac2ca74..6aca356bf5 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -23,14 +23,15 @@ # pylint: disable=wrong-import-position # This import needs to be after pytest in order to prevent ImportErrors import pennylane as qml +from xdsl.dialects import test +from xdsl.dialects.builtin import ModuleOp +from xdsl.ir.core import Block, Region + from catalyst.python_interface.conversion import xdsl_from_qjit from catalyst.python_interface.visualization.construct_circuit_dag import ( ConstructCircuitDAG, ) from catalyst.python_interface.visualization.dag_builder import DAGBuilder -from xdsl.dialects import test -from xdsl.dialects.builtin import ModuleOp -from xdsl.ir.core import Block, Region class FakeDAGBuilder(DAGBuilder): @@ -491,7 +492,7 @@ def test_nested_conditionals(self): @xdsl_from_qjit @qml.qjit(autograph=True, target="mlir") @qml.qnode(dev) - def my_workflow(x,y): + def my_workflow(x, y): if x == 1: if y == 2: qml.H(0) @@ -501,7 +502,7 @@ def my_workflow(x,y): else: qml.Z(0) - args = (1,2) + args = (1, 2) module = my_workflow(*args) utility = ConstructCircuitDAG(FakeDAGBuilder()) From 9996684c531a94b6d116a57b4519f073ceeea243 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 13:20:15 -0500 Subject: [PATCH 237/245] fix test --- .../test_construct_circuit_dag.py | 24 +++++++++++++------ 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 6aca356bf5..4834005766 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -513,20 +513,30 @@ def my_workflow(x, y): # cluster0 -> qjit # cluster1 -> my_workflow + # cluster2 -> conditional (1) + # cluster3 -> if + # cluster4 -> conditional () + # cluster5 -> if + # cluster6 -> else + # cluster7 -> else # Check first conditional is a cluster within my_workflow assert clusters["cluster2"]["cluster_label"] == "conditional" assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" - # Check two clusters live within first conditional + # Check 'if' cluster of first conditional has another conditional assert clusters["cluster3"]["node_label"] == "if ..." assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" - # first conditional's else - assert clusters["cluster6"]["node_label"] == "else" - assert clusters["cluster6"]["parent_cluster_uid"] == "cluster2" - # Check nested if / else is within the first if cluster - assert clusters["cluster4"]["node_label"] == "if ..." + # Second conditional + assert clusters["cluster4"]["cluster_label"] == "conditional" assert clusters["cluster4"]["parent_cluster_uid"] == "cluster3" + # Check 'if' and 'else' in second conditional assert clusters["cluster5"]["node_label"] == "if ..." - assert clusters["cluster5"]["parent_cluster_uid"] == "cluster3" + assert clusters["cluster5"]["parent_cluster_uid"] == "cluster4" + assert clusters["cluster6"]["node_label"] == "else" + assert clusters["cluster6"]["parent_cluster_uid"] == "cluster4" + + # Check nested if / else is within the first if cluster + assert clusters["cluster7"]["node_label"] == "else" + assert clusters["cluster7"]["parent_cluster_uid"] == "cluster2" From c262df9dcce420b083241cf8dca70f0dadd8cda6 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 14:11:14 -0500 Subject: [PATCH 238/245] format --- .../visualization/test_construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 4834005766..48005befe4 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -516,8 +516,8 @@ def my_workflow(x, y): # cluster2 -> conditional (1) # cluster3 -> if # cluster4 -> conditional () - # cluster5 -> if - # cluster6 -> else + # cluster5 -> if + # cluster6 -> else # cluster7 -> else # Check first conditional is a cluster within my_workflow From 506f8398f3a7a0dcea93eeefa0f1c421982c30fc Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 17:04:09 -0500 Subject: [PATCH 239/245] fix labels --- .../visualization/construct_circuit_dag.py | 13 ++++------- .../test_construct_circuit_dag.py | 22 +++++++++---------- 2 files changed, 15 insertions(+), 20 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 4f63fe26a4..4919a56e0f 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -94,15 +94,10 @@ def _visit_block(self, block: Block) -> None: def _for_op(self, operation: scf.ForOp) -> None: """Handle an xDSL ForOp operation.""" - # TODO: Extract from IR in future PR - iter_var = "..." - start, stop, step = "...", "...", "..." - label = f"for {iter_var} in range({start}, {stop}, {step})" - uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, - node_label=label, + node_label="for loop", label="", cluster_uid=self._cluster_uid_stack[-1], ) @@ -120,7 +115,7 @@ def _while_op(self, operation: scf.WhileOp) -> None: uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( uid, - node_label="while ...", + node_label="while loop", label="", cluster_uid=self._cluster_uid_stack[-1], ) @@ -155,11 +150,11 @@ def _if_op(self, operation: scf.IfOp): def _get_conditional_branch_label(i): if i == 0: - return "if ..." + return "if" elif i == num_regions - 1: return "else" else: - return "elif ..." + return "elif" uid = f"cluster{self._cluster_uid_counter}" self.dag_builder.add_cluster( diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 48005befe4..07bca62cba 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -314,7 +314,7 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters["cluster2"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster2"]["node_label"] == "for loop" assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" @pytest.mark.unit @@ -340,9 +340,9 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters["cluster2"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster2"]["node_label"] == "for loop" assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" - assert clusters["cluster3"]["node_label"] == "for ... in range(..., ..., ...)" + assert clusters["cluster3"]["node_label"] == "for loop" assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" @@ -372,7 +372,7 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters["cluster2"]["node_label"] == "while ..." + assert clusters["cluster2"]["node_label"] == "while loop" assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" @pytest.mark.unit @@ -402,9 +402,9 @@ def my_workflow(): # cluster0 -> qjit # cluster1 -> my_workflow - assert clusters["cluster2"]["node_label"] == "while ..." + assert clusters["cluster2"]["node_label"] == "while loop" assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" - assert clusters["cluster3"]["node_label"] == "while ..." + assert clusters["cluster3"]["node_label"] == "while loop" assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" @@ -440,7 +440,7 @@ def my_workflow(x): assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" # Check three clusters live within cluster2 (conditional) - assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["node_label"] == "if" assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" assert clusters["cluster4"]["node_label"] == "else" assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" @@ -476,9 +476,9 @@ def my_workflow(x): assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" # Check three clusters live within conditional - assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["node_label"] == "if" assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" - assert clusters["cluster4"]["node_label"] == "elif ..." + assert clusters["cluster4"]["node_label"] == "elif" assert clusters["cluster4"]["parent_cluster_uid"] == "cluster2" assert clusters["cluster5"]["node_label"] == "else" assert clusters["cluster5"]["parent_cluster_uid"] == "cluster2" @@ -525,14 +525,14 @@ def my_workflow(x, y): assert clusters["cluster2"]["parent_cluster_uid"] == "cluster1" # Check 'if' cluster of first conditional has another conditional - assert clusters["cluster3"]["node_label"] == "if ..." + assert clusters["cluster3"]["node_label"] == "if" assert clusters["cluster3"]["parent_cluster_uid"] == "cluster2" # Second conditional assert clusters["cluster4"]["cluster_label"] == "conditional" assert clusters["cluster4"]["parent_cluster_uid"] == "cluster3" # Check 'if' and 'else' in second conditional - assert clusters["cluster5"]["node_label"] == "if ..." + assert clusters["cluster5"]["node_label"] == "if" assert clusters["cluster5"]["parent_cluster_uid"] == "cluster4" assert clusters["cluster6"]["node_label"] == "else" assert clusters["cluster6"]["parent_cluster_uid"] == "cluster4" From f67802fa6130614c35f74ac0f3fd093a7a6a2d05 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Thu, 4 Dec 2025 17:07:02 -0500 Subject: [PATCH 240/245] fix labels --- .../python_interface/visualization/construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 4919a56e0f..55efb1bd78 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -137,7 +137,6 @@ def _if_op(self, operation: scf.IfOp): uid, node_label="", label="conditional", - fontsize=10, labeljust="l", cluster_uid=self._cluster_uid_stack[-1], ) From d069fef466ace630443a11a964d38e9a17001066 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 5 Dec 2025 12:08:33 -0500 Subject: [PATCH 241/245] fix --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 237fc3c469..5181abd2d4 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -19,7 +19,7 @@ from xdsl.dialects import builtin, func, scf from xdsl.ir import Block, Operation, Region -from catalyst.python_interface.dialects import catalyst, quantum +from catalyst.python_interface.dialects import quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder From e29f5468f4f5dccaa99c6a66bf2fdf8e445fd5b0 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 5 Dec 2025 12:09:52 -0500 Subject: [PATCH 242/245] bring back ssavalue --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 5181abd2d4..6ded0d80ae 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,7 +16,7 @@ from functools import singledispatchmethod -from xdsl.dialects import builtin, func, scf +from xdsl.dialects import builtin, func, scf, SSAValue from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import quantum From b34d8e2b0c590bd9c08bfa1928b4ab2fec4b5d17 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 5 Dec 2025 12:10:03 -0500 Subject: [PATCH 243/245] format --- .../python_interface/visualization/construct_circuit_dag.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 6ded0d80ae..44b0d5807e 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,7 +16,7 @@ from functools import singledispatchmethod -from xdsl.dialects import builtin, func, scf, SSAValue +from xdsl.dialects import SSAValue, builtin, func, scf from xdsl.ir import Block, Operation, Region from catalyst.python_interface.dialects import quantum From aa1d90ce1b61c097b33596a81a780cbd46edcec4 Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 5 Dec 2025 12:10:50 -0500 Subject: [PATCH 244/245] format --- .../python_interface/visualization/construct_circuit_dag.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py index 44b0d5807e..6a548b5d51 100644 --- a/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py +++ b/frontend/catalyst/python_interface/visualization/construct_circuit_dag.py @@ -16,8 +16,8 @@ from functools import singledispatchmethod -from xdsl.dialects import SSAValue, builtin, func, scf -from xdsl.ir import Block, Operation, Region +from xdsl.dialects import builtin, func, scf +from xdsl.ir import Block, Operation, Region, SSAValue from catalyst.python_interface.dialects import quantum from catalyst.python_interface.visualization.dag_builder import DAGBuilder From 27ff971f4caea90009b646ad4a7f0c86105de39f Mon Sep 17 00:00:00 2001 From: andrijapau Date: Fri, 5 Dec 2025 12:12:33 -0500 Subject: [PATCH 245/245] remove random import --- .../python_interface/visualization/test_construct_circuit_dag.py | 1 - 1 file changed, 1 deletion(-) diff --git a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py index 3ee18263b4..82f50140fa 100644 --- a/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py +++ b/frontend/test/pytest/python_interface/visualization/test_construct_circuit_dag.py @@ -16,7 +16,6 @@ from unittest.mock import Mock import pytest -from jax import util pytestmark = pytest.mark.usefixtures("requires_xdsl")