Fix qml.prod with autograph

[rniczh]

Context:

Fix qml.prod with autograph

@qml.prod
def template(b):
    if b:
        qml.H(0)
        qml.X(0)

@qjit(autograph=True)
@qml.qnode(qml.device("null.qubit", wires=1))
def circuit(b: bool):
    template(b)
    return qml.state()

print(circuit(True))

Error message:

jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function circuit at ... for qjit_capture. This concrete value was not available in Python because it depends on the value of the argument b.
See https://docs.jax.dev/en/latest/errors.html#jax.errors.TracerBoolConversionError

Description of the Change:

Handle calls to functions that were decorated with qml.prod, qml.adjoint, etc. These decorators return wrapper functions that call the original function without autograph conversion. We detect these wrappers and unwrap them to convert the original function with autograph.

Benefits:

Possible Drawbacks:

Related GitHub Issues:
[sc-95653]
#1911

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 97.49%. Comparing base (f0c21f5) to head (67505c5).
⚠️ Report is 1 commits behind head on main.

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #2083   +/-   ##
=======================================
  Coverage   97.48%   97.49%           
=======================================
  Files          91       91           
  Lines       10594    10601    +7     
  Branches      990      992    +2     
=======================================
+ Hits        10328    10335    +7     
  Misses        211      211           
  Partials       55       55           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[github-actions]

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.

[dime10]

Just wondering, are any of the changes besides this line relevant, or just refactoring?

[rniczh]

Yes, it's because I use it to find the decorator and apply it again to the converted call of the fn.__wrapped.
Actually, there is another way to achieve the same behaviour is through fn.__globals__ (as shown below). It doesn't need to create the _known_wrapper_functions, and also doesn't need to find which decorator need to be applied as well. With that being said, it didn't seem robust enough for me, so I didn't use this solution 🤔. But if you find this solution is better than the previous one, I'm happy to update it 👍

decorator_name = fn.__module__.split(".")[-1]
decorator = fn.__globals__.get(decorator_name)

catalyst/frontend/catalyst/autograph/ag_primitives.py

Lines 611 to 613 in 67505c5

	if decorator := next(
	(f for f in _known_wrapper_functions if f.__module__ == fn.__module__), None
	):

[dime10]

Yes, it's because I use it to find the decorator and apply it again to the converted call of the fn.__wrapped

Right, I am wondering about the robustness, but maybe going via wrapped is actually the best way here? Not sure.

Normally, we grab the decorated function here since it's the first argument to the decorator.
The exception to this is if the decorator was already applied, outside the qjit scope, in which case it's too late to hit this branch. For that we use the CatalystCallable, or handle the QNode specially. I guess the difference with qml.prod and the other qml decorators here is that ctrl and adjoint are dispatched to their catalyst equivalent, so can be handled via CatalystCallable, whereas qml.prod is not dispatched 🤔

[rniczh]

I guess the difference with qml.prod and the other qml decorators here is that ctrl and adjoint are dispatched to their catalyst equivalent, so can be handled via CatalystCallable, whereas qml.prod is not dispatched

Yes, exactly, and prod aslo doesn't has the catalyst equivalent impl as well.