PennyLaneAI / PennyLaneAI/catalyst

[Frontend] Result type unification across conditional branches might not work for qnode-programs

Open
#824 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
234
Forks
84
Avg merge
2d 15h
Merged PRs (30d)
66

Description

This issue reflexes the fixme comment in the quantum tracing pipeline for the conditional primitive.

The core of the problem is: we re-trace Jaxpr program in order to unify the result values (link) but at the same time, for quantum programs, we use two-staged classical-quantum tracing. Before quantum tracing is complete, the Jaxpr IR might not be valid from the Jax point of view. Some tracers might not have their source equations yet, e.g. this might happen with programs involving qml.measure. Because of this, we only perform the unification during the quantum tracing (link), which might be too late.

In order to fix this problem, we might need to either (1) guarantee that the IR is always valid after the classical tracing is complete which might be hard or (2) Implement the unification "manually" in a way that would allow dangling Jaxpr variables.

The second approach seems be more realistic in the short-term.

See also:

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with the fixme in frontend/catalyst/api_extensions/control_flow.py and the retracing logic in frontend/catalyst/jax_tracer.py. Compare the xfailed case at frontend/test/pytest/test_conditionals.py:313 with the passing pure-qjit case at line 294, then verify that result unification works for qnode-program conditionals without invalid Jaxpr assumptions.

Written by the indexing model from the issue text.

Assessment

Tech stack
python
Domain
compilers
Issue type
Bug
Difficulty
5/5
Estimated time
Over a week
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.