proof of concept for capturing constants inside qnode

Author: erick-xanadu

frontend/catalyst/api_extensions/differentiation.py

@@ -888,16 +888,3 @@ def _check_qnode_against_grad_method(f: QNode, method: str, jaxpr: Jaxpr):
             "Cannot differentiate a QNode explicitly marked non-differentiable (with "
             "diff_method=None)."
         )
-
-    if f.diff_method == "parameter-shift" and any(
-        prim not in [expval_p, probs_p] for prim in return_ops
-    ):
-        raise DifferentiableCompileError(
-            "The parameter-shift method can only be used for QNodes "
-            "which return either qml.expval or qml.probs."
-        )
-
-    if f.diff_method == "adjoint" and any(prim not in [expval_p] for prim in return_ops):
-        raise DifferentiableCompileError(
-            "The adjoint method can only be used for QNodes which return qml.expval."
-        )

frontend/catalyst/pipelines.py

@@ -173,6 +173,7 @@ def get_enforce_runtime_invariants_stage(_options: CompileOptions) -> List[str]:
         # keep inlining modules targeting the Catalyst runtime.
         # But qnodes targeting other backends may choose to lower
         # this into something else.
+        "builtin.module(inline)",
         "inline-nested-module",
     ]
     return enforce_runtime_invariants
@@ -217,7 +218,6 @@ def get_bufferization_stage(_options: CompileOptions) -> List[str]:
     """Returns the list of passes that performs bufferization"""
     bufferization = [
         "one-shot-bufferize{dialect-filter=memref}",
-        "inline",
         "gradient-preprocess",
         "gradient-bufferize",
         "scf-bufferize",

frontend/catalyst/qfunc.py

@@ -131,6 +131,9 @@ def __call__(self, *args, **kwargs):
         out_tree_expected = kwargs.pop("_out_tree_expected", [])
         debug_info = kwargs.pop("debug_info", None)
 
+        import jax
+
+        @jax.jit
         def _eval_quantum(*args, **kwargs):
             closed_jaxpr, out_type, out_tree, out_tree_exp = trace_quantum_function(
                 self.func,