Enable Autograph on qml functions (#1212)

Fix a bug preventing the target of `qml.adjoint` and `qml.ctrl` calls
from being transformed by AutoGraph.

doc/releases/changelog-dev.md

@@ -273,6 +273,9 @@
 
 <h3>Bug fixes</h3>
 
+* Fix a bug preventing the target of `qml.adjoint` and `qml.ctrl` calls from being transformed by
+  AutoGraph.
+  [(#1212)](https://github.com/PennyLaneAI/catalyst/pull/1212)
 
 * Resolve a bug where `mitigate_with_zne` does not work properly with shots and devices
   supporting only Counts and Samples (e.g. Qrack). (transform: `measurements_from_sample`).

frontend/catalyst/autograph/ag_primitives.py

@@ -537,7 +537,9 @@ def converted_call(fn, args, kwargs, caller_fn_scope=None, options=None):
         # HOTFIX: pass through calls of known Catalyst wrapper functions
         if fn in (
             catalyst.adjoint,
+            qml.adjoint,
             catalyst.ctrl,
+            qml.ctrl,
             catalyst.grad,
             catalyst.value_and_grad,
             catalyst.jacobian,

frontend/test/pytest/test_autograph.py

@@ -25,25 +25,7 @@
 from jax.errors import TracerBoolConversionError
 from numpy.testing import assert_allclose
 
-from catalyst import (
-    AutoGraphError,
-    adjoint,
-    autograph_source,
-    cond,
-    ctrl,
-    debug,
-    disable_autograph,
-    for_loop,
-    grad,
-    jacobian,
-    jvp,
-    measure,
-    qjit,
-    run_autograph,
-    vjp,
-    vmap,
-    while_loop,
-)
+from catalyst import *
 from catalyst.autograph.transformer import TRANSFORMER
 from catalyst.utils.dummy import dummy_func
 from catalyst.utils.exceptions import CompileError
@@ -295,7 +277,8 @@ def fn(x: float):
         assert check_cache(inner.user_function.func)
         assert fn(np.pi) == -1
 
-    def test_adjoint_wrapper(self):
+    @pytest.mark.parametrize("adjoint_fn", [adjoint, qml.adjoint])
+    def test_adjoint_wrapper(self, adjoint_fn):
         """Test conversion is happening succesfully on functions wrapped with 'adjoint'."""
 
         def inner(x):
@@ -304,14 +287,15 @@ def inner(x):
         @qjit(autograph=True)
         @qml.qnode(qml.device("lightning.qubit", wires=1))
         def fn(x: float):
-            adjoint(inner)(x)
+            adjoint_fn(inner)(x)
             return qml.probs()
 
         assert hasattr(fn.user_function, "ag_unconverted")
         assert check_cache(inner)
         assert np.allclose(fn(np.pi), [0.0, 1.0])
 
-    def test_ctrl_wrapper(self):
+    @pytest.mark.parametrize("ctrl_fn", [ctrl, qml.ctrl])
+    def test_ctrl_wrapper(self, ctrl_fn):
         """Test conversion is happening succesfully on functions wrapped with 'ctrl'."""
 
         def inner(x):
@@ -320,7 +304,7 @@ def inner(x):
         @qjit(autograph=True)
         @qml.qnode(qml.device("lightning.qubit", wires=2))
         def fn(x: float):
-            ctrl(inner, control=1)(x)
+            ctrl_fn(inner, control=1)(x)
             return qml.probs()
 
         assert hasattr(fn.user_function, "ag_unconverted")