Integrate vjp with program capture

.dep-versions

@@ -10,7 +10,7 @@ enzyme=v0.0.203
 
 # For a custom PL version, update the package version here and at
 # 'doc/requirements.txt'
-pennylane=0.45.0-dev12
+pennylane=0.45.0-dev13
 
 # For a custom LQ/LK version, update the package version here and at
 # 'doc/requirements.txt'

doc/releases/changelog-dev.md

@@ -17,8 +17,15 @@
   length and the number of qubit operands are the same, and that all of the Pauli words are legal.
   [(#2405)](https://github.com/PennyLaneAI/catalyst/pull/2405)
 
+* `qml.vjp` can now be used with Catalyst and program capture.
+  [(#2279)](https://github.com/PennyLaneAI/catalyst/pull/2279)
+
 <h3>Breaking changes 💔</h3>
 
+* When an integer argnums is provided to `catalyst.vjp`, a singleton dimension is now squeezed
+  out. This brings the behaviour in line with that of `grad` and `jacobian`.
+  [(#2279)](https://github.com/PennyLaneAI/catalyst/pull/2279)
+
 * Dropped support for NumPy 1.x following its end-of-life. NumPy 2.0 or higher is now required.
   [(#2407)](https://github.com/PennyLaneAI/catalyst/pull/2407)
 

doc/requirements.txt

@@ -34,4 +34,4 @@ lxml_html_clean
 --extra-index-url https://test.pypi.org/simple/
 pennylane-lightning-kokkos==0.45.0-dev8
 pennylane-lightning==0.45.0-dev8
-pennylane==0.45.0-dev12
+pennylane==0.45.0-dev13

frontend/catalyst/api_extensions/differentiation.py

@@ -24,6 +24,7 @@
 from typing import Callable, Iterable, List, Optional, Union
 
 import jax
+import pennylane as qml
 from jax._src.api import _dtype
 from jax._src.tree_util import PyTreeDef, tree_flatten, tree_unflatten
 from jax.api_util import debug_info
@@ -546,13 +547,14 @@ def f(x):
     (Array([0.09983342, 0.04      , 0.02      ], dtype=float64),
      (Array([-0.43750208,  0.07      ], dtype=float64),))
     """
+    if qml.capture.enabled():
+        return qml.vjp(f, params, cotangents, method=method, h=h, argnums=argnums)
 
     def check_is_iterable(x, hint):
         if not isinstance(x, Iterable):
             raise ValueError(f"vjp '{hint}' argument must be an iterable, not {type(x)}")
 
     check_is_iterable(params, "params")
-    check_is_iterable(cotangents, "cotangents")
 
     if EvaluationContext.is_tracing():
         scalar_out = False
@@ -564,7 +566,10 @@ def check_is_iterable(x, hint):
         grad_params = _check_grad_params(method, scalar_out, h, argnums, len(args_flatten), in_tree)
 
         args_argnums = tuple(params[i] for i in grad_params.argnums)
-        _, in_tree = tree_flatten(args_argnums)
+        if isinstance(argnums, int) or argnums is None:
+            _, in_tree = tree_flatten(0)
+        else:
+            _, in_tree = tree_flatten(args_argnums)
 
         jaxpr, out_tree = _make_jaxpr_check_differentiable(fn, grad_params, *params)
 

frontend/catalyst/autograph/ag_primitives.py

@@ -532,6 +532,10 @@ def converted_call(fn, args, kwargs, caller_fn_scope=None, options=None):
             qml.prod,
             catalyst.ctrl,
             qml.ctrl,
+            qml.grad,
+            qml.jacobian,
+            qml.vjp,
+            qml.jvp,
             catalyst.grad,
             catalyst.value_and_grad,
             catalyst.jacobian,

frontend/catalyst/from_plxpr/from_plxpr.py

@@ -29,6 +29,7 @@
 from pennylane.capture.expand_transforms import ExpandTransformsInterpreter
 from pennylane.capture.primitives import jacobian_prim as pl_jac_prim
 from pennylane.capture.primitives import transform_prim
+from pennylane.capture.primitives import vjp_prim as pl_vjp_prim
 from pennylane.transforms import commute_controlled as pl_commute_controlled
 from pennylane.transforms import decompose as pl_decompose
 from pennylane.transforms import gridsynth as pl_gridsynth
@@ -239,6 +240,18 @@ def handle_grad(self, *args, jaxpr, n_consts, **kwargs):
     )
 
 
+@WorkflowInterpreter.register_primitive(pl_vjp_prim)
+def handle_vjp(self, *args, jaxpr, **kwargs):
+    """Translate a grad equation."""
+    f = partial(copy(self).eval, jaxpr, [])
+    new_jaxpr = jax.make_jaxpr(f)(*args[: -len(jaxpr.outvars)])
+
+    new_args = (*new_jaxpr.consts, *args)
+    j = new_jaxpr.jaxpr
+    new_j = j.replace(constvars=(), invars=j.constvars + j.invars)
+    return pl_vjp_prim.bind(*new_args, jaxpr=new_j, **kwargs)
+
+
 # pylint: disable=unused-argument, too-many-arguments
 @WorkflowInterpreter.register_primitive(qnode_prim)
 def handle_qnode(

frontend/catalyst/jax_primitives.py

@@ -134,6 +134,7 @@
     )
 
 from pennylane.capture.primitives import jacobian_prim as pl_jac_prim
+from pennylane.capture.primitives import vjp_prim as pl_vjp_prim
 
 from catalyst.compiler import get_lib_path
 from catalyst.jax_extras import (
@@ -931,6 +932,38 @@ def _vjp_lowering(ctx, *args, jaxpr, fn, grad_params):
     ).results
 
 
+def _capture_vjp_lowering(ctx, *args, jaxpr, fn, method, argnums, h):
+    """
+    Returns:
+        MLIR results
+    """
+    args = list(args)
+    mlir_ctx = ctx.module_context.context
+    n_params = len(jaxpr.invars)
+    new_argnums = np.array(argnums)
+
+    output_types = list(map(mlir.aval_to_ir_types, ctx.avals_out))
+    flat_output_types = util.flatten(output_types)
+    func_args = args[:n_params]
+    cotang_args = args[n_params:]
+    func_result_types = flat_output_types[: len(flat_output_types) - len(argnums)]
+    vjp_result_types = flat_output_types[len(flat_output_types) - len(argnums) :]
+
+    func_op = lower_jaxpr(ctx, jaxpr, (method, h, *argnums), fn=fn)
+
+    symbol_ref = get_symbolref(ctx, func_op)
+    return VJPOp(
+        func_result_types,
+        vjp_result_types,
+        ir.StringAttr.get(method),
+        symbol_ref,
+        mlir.flatten_ir_values(func_args),
+        mlir.flatten_ir_values(cotang_args),
+        diffArgIndices=ir.DenseIntElementsAttr.get(new_argnums),
+        finiteDiffParam=ir.FloatAttr.get(ir.F64Type.get(mlir_ctx), h) if h else None,
+    ).results
+
+
 #
 # zne
 #
@@ -2843,6 +2876,7 @@ def subroutine_lowering(*args, **kwargs):
     (for_p, _for_loop_lowering),
     (grad_p, _grad_lowering),
     (pl_jac_prim, _capture_grad_lowering),
+    (pl_vjp_prim, _capture_vjp_lowering),
     (func_p, _func_lowering),
     (jvp_p, _jvp_lowering),
     (vjp_p, _vjp_lowering),

frontend/test/pytest/test_autograph.py

@@ -396,19 +396,26 @@ def fn(x: float):
         assert check_cache(inner)
         assert fn(3) == tuple([jax.numpy.array(2.0), jax.numpy.array(6.0)])
 
+    @pytest.mark.usefixtures("use_both_frontend")
     @pytest.mark.parametrize("vjp_func", [vjp, qml.vjp])
     def test_vjp_wrapper(self, vjp_func):
         """Test conversion is happening succesfully on functions wrapped with 'vjp'."""
 
+        if qml.capture.enabled() and vjp_func == vjp:  # pylint: disable=comparison-with-callable
+            pytest.xfail("program capture autograph doesn't work with catalyst.vjp")
+
         def inner(x):
-            return 2 * x, x**2
+            if x > 0:
+                return 2 * x, x**2
+            return 4 * x, x**8
 
         @qjit(autograph=True)
         def fn(x: float):
             return vjp_func(inner, (x,), (1.0, 1.0))
 
         assert hasattr(fn.user_function, "ag_unconverted")
-        assert check_cache(inner)
+        if not qml.capture.enabled():
+            assert check_cache(inner)
         assert np.allclose(fn(3)[0], tuple([jnp.array(6.0), jnp.array(9.0)]))
         assert np.allclose(fn(3)[1], jnp.array(8.0))