Graph decomposition for device and stopping_condition by lazypanda10117 · Pull Request #2472 · PennyLaneAI/catalyst

lazypanda10117 · 2026-02-10T15:35:04Z

Context:

Supporting device-specific gateset decomposition in the program capture flow. This is enabled by extending the handle_qnode function. As stated in Ali's comment, the new handle_qnode supports the following cases:

User defined stopping_condition in qml.decompose: Use PLxPR decompose with user stopping condition!
No user decomposition (default): Load device capabilities --> device-specific gate set --> Use graph and decompose-lowering. Fallback to PLxPR when graph fails.
User specified gate_set : Graph-based decomposition on user gate set (the current support). Fallback to PLxPR when graph fails.

[sc-110277]
[sc-110276]
[sc-107559]

github-actions · 2026-02-10T22:38:15Z

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md on your branch with:

A one-to-two sentence description of the change. You may include a small working example for new features.
A link back to this PR.
Your name (or GitHub username) in the contributors section.

maliasadi

Thanks @lazypanda10117! The main remaining TODO is to report/fix the edge cases:

Counts Not Implemented with Capture <- this issue is a program capture limitation and doesn't have anything to do with decomposition
Almost all xDSL / PPR parse errors are related to producing a different IR as of applying decomposition! We can fix most of them by updating the MLIR output to be compatible with the checks.
HybridOp decomp assertion <-- we should be able to fix all these issues by ensureing that HybridOps are in the stopping _condition OR excluded from decomposition.
Ops missing from device gateset <- Lightning/null devices support more gates than Catalyst can represent! You need to take the intersection of these ops and decompose the remaining ones in PLxPR.
- KeyError: 'StatePrep'
- KeyError: 'BasisState'
- KeyError: 'QubitUnitary'
- KeyError: 'QFT'
_OperatorNode Can't Be Captured Without Wire Count <- the graph decomposition produces _OperatorNode objects that require a wire count to be captured/lowered into MLIR, but the wire count isn't available at graph construction time. cc: @astralcai
In specs/pass counting/circuit drawing tests: the current pipeline adds extra decompose-lowering pass to the compiled module. We need to update the logic (and maybe tests) to avoid this.
there are some other errors such as DID NOT RAISE/WARN or Cannot resolve wire for SliceOp or TracerBoolConversionError; let's get back to them after fixing the above as some of them are tied with those edge cases.

maliasadi · 2026-02-11T18:42:10Z

    consts = args[shots_len : n_consts + shots_len]
    non_const_args = args[shots_len + n_consts :]

-    closed_jaxpr = (


The new handle_qnode supports the following cases:

User defined stopping_condition in qml.decompose: Use PLxPR decompose with user stopping condition!

No user decomposition (default): Load device capabilities --> device-specific gate set --> Use graph and decompose-lowering

User specified gate_set : Graph-based decomposition on user gate set (the current support)

cc: @astralcai @isaacdevlugt if you guys have any comments or suggestions

If a user applies a @qml.decompose on top of a qnode, it would be applied in addition to the decomposition to the device gate set right? That is the current behaviour at least. Whether the user has a @qml.decompose or not does not change the fact that a decomposition to the device gate set is always going to happen. Is that not the case here?

If a user applies a @qml.decompose on top of a qnode, it would be applied in addition to the decomposition to the device gate set right?

My understanding is that when a user specifies a transform, we will only use that user-specified gateset. We only do device-level decomposition when no such user-defined transform is provided.

maliasadi · 2026-02-11T18:47:47Z

+
+def calculate_diff_method(device, closed_jaxpr):
+    """Calculate the diff method for the device."""
+    return _calculate_diff_method(device, closed_jaxpr)


is this correct? 🤔

We don't have any logic to deal with "diff" ops in the decomposition system

This is used in catalyst acceptance. I am just porting over the function.

maliasadi · 2026-02-11T18:50:33Z

Could we merge the util methods with the one in ./catalyst/device or add these methods to from_plxpr/decompose.py ?

Yes we can, I just think this is cleaner since this is the subset of functions (with slight modifications) that are needed from device capabilities.

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

lazypanda10117 · 2026-02-11T20:40:35Z

Ops missing from device gateset <- Lightning/null devices support more gates than Catalyst can represent! You need to take the intersection of these ops and decompose the remaining ones in PLxPR.

I think the issue with (4) is actually due to apply_compiler_decompose_to_plxpr. In that function, we do disable_graph and then enable_graph. However, sometimes, graph is not necessarily enabled when we initially called _apply_compiler_decompose_to_plxpr. When running parameterized tests, the environment is not reset across each parameterized run, so this makes subsequent tests run with graph enabled, causing the failures above.

astralcai · 2026-02-11T21:22:14Z

_OperatorNode Can't Be Captured Without Wire Count <- the graph decomposition produces _OperatorNode objects that require a wire count to be captured/lowered into MLIR, but the wire count isn't available at graph construction time. cc: @astralcai

As discussed offline with @lazypanda10117, the issue with this type of errors comes from operators like RotXZX not supported in Catalyst. We need to do a first pass decomposition into the compiler gate set.

lazypanda10117 · 2026-02-12T18:17:45Z

@maliasadi Point 3,4,5,6 are mostly addressed.

**Context:** Fix failing test after PennyLaneAI/pennylane#9001 **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:**

**Context:** We are about to start the audit of tests. For convenience, a tiny nice fixture to automatically apply the two values of `capture` kwarg locally is good. **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:** [sc-111330] --------- Co-authored-by: Andrija Paurevic <46359773+andrijapau@users.noreply.github.com>

**Context:** The PPR/PPM lowering passes (`lower-qec-init-ops`, `unroll-conditional-ppr-ppm`) are now run as part of the main quantum compilation pipeline. When using `to-ppr` and `ppr-to-ppm` transforms, these passes are applied automatically during compilation; we no longer need to stack them explicitly. [[sc-109665]] --------- Co-authored-by: Paul <79805239+paul0403@users.noreply.github.com> Co-authored-by: Hong-Sheng Zheng <mathan0203@gmail.com>

…h-decomp-stop-condition

Context: Supporting `stopping_condition` when user supplies it in `qml.decompose` with capture enabled. This will use the plxpr transform to handle the decomposition with stopping condition. This PR is a subset of #2472. See that PR for a more detailed discussion. [[sc-110276]] --------- Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

lazypanda10117 added 8 commits January 29, 2026 17:01

first version with stopping condition

70a4b0c

Decomposition with device

570b689

Refactoring

43eaa7b

Removing comments and prints

f9c726c

Formatting

4b53bb2

Different decomposition path

922a14a

Format

eee59bc

Merge branch 'main' into lazypanda10117/graph-decomp-stop-condition

6b21d65

lazypanda10117 marked this pull request as ready for review February 10, 2026 22:37

lazypanda10117 requested a review from maliasadi February 10, 2026 22:38

lazypanda10117 added the wip PRs that are a Work-In-Progress label Feb 10, 2026

lazypanda10117 added 2 commits February 10, 2026 17:46

fixing gateset intersection

8e0d3b7

Fixing some bugs

1b7acf9

maliasadi requested changes Feb 11, 2026

View reviewed changes

Apply suggestion from @maliasadi

01373c0

Co-authored-by: Ali Asadi <10773383+maliasadi@users.noreply.github.com>

lazypanda10117 added 2 commits February 11, 2026 18:02

Fixing decomposition logic

b0b9a37

Fixing PPR tests

75a4d82

maliasadi reviewed Feb 12, 2026

View reviewed changes

Comment thread frontend/catalyst/device/qjit_device.py Outdated

lazypanda10117 added 2 commits February 12, 2026 13:13

Removing changes in device

6084f93

update

b9183ee

ringo-but-quantum and others added 6 commits February 12, 2026 16:55

[no ci] bump nightly version

fd345ff

Fix expected warning type in test (#2477)

76dbb40

**Context:** Fix failing test after PennyLaneAI/pennylane#9001 **Description of the Change:** **Benefits:** **Possible Drawbacks:** **Related GitHub Issues:**

Add libmlir_apfloat_wrappers to Catalyst wheels (#2478)

2165da8

[no ci] bump nightly version

4019759

lazypanda10117 added 19 commits February 17, 2026 09:57

Merge branch 'main' into lazypanda10117/graph-decomp-stop-condition

68383e9

Upddating dep-version

7d3da80

Fixing pretty print issue

d02b711

Fixing string mismatch issue

ec3abdb

Fixing kwargs name typo

2ec5c3a

Fixing many issues

07cc7b4

Adding back the templated decomposition for graph decomp

ca5a678

Pulling changes from lazypanda10117/graph-decomp-stop-condition

0ff1806

Removing extraneous changes

ed465e1

QEC to PBC

c3f174f

Merge branch 'main' into lazypanda10117/graph-decomp-stop-condition

ea9398f

Fixing Jax interp issues

4729be2

Fixing failing tests

221f1e7

Fixing all errors

28c267d

removing prints

b123644

Merge branch 'main' into lazypanda10117/stop-condition

c1abf8c

format and docs

b36bc05

Adding a test for stopping condition

3e9eafd

Formatting

dbcf495

lazypanda10117 mentioned this pull request Feb 18, 2026

Supporting Stopping Condition with qml.decompose #2486

Merged

Changelog

c87473c

lazypanda10117 changed the base branch from main to lazypanda10117/stop-condition February 18, 2026 21:58

lazypanda10117 added 5 commits February 18, 2026 17:08

Removing print

b1a8db0

Fixing specs failing test case

48c5eb3

Format

00e14c9

Merge branch 'lazypanda10117/stop-condition' into lazypanda10117/grap…

e5b9567

…h-decomp-stop-condition

Adding the use_graph variable

54a3b6e

Base automatically changed from lazypanda10117/stop-condition to main February 23, 2026 19:53

previous changes

1023a3d

Conversation

lazypanda10117 commented Feb 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Context:

Uh oh!

github-actions Bot commented Feb 10, 2026

Uh oh!

maliasadi left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

maliasadi Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

astralcai Feb 13, 2026

Choose a reason for hiding this comment

Uh oh!

lazypanda10117 Feb 17, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

maliasadi Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

lazypanda10117 Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

maliasadi Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

lazypanda10117 Feb 11, 2026

Choose a reason for hiding this comment

Uh oh!

lazypanda10117 commented Feb 11, 2026

Uh oh!

astralcai commented Feb 11, 2026

Uh oh!

Uh oh!

lazypanda10117 commented Feb 12, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

lazypanda10117 commented Feb 10, 2026 •

edited

Loading

maliasadi left a comment •

edited

Loading