Fix compilation of Jacobian with for loops and vmap

doc/releases/changelog-dev.md

@@ -44,6 +44,25 @@
 
 <h3>Bug fixes</h3>
 
+* Circuits with preprocessing functions outside qnodes can now be differentiated.
+  [(#332)](https://github.com/PennyLaneAI/catalyst/pull/332)
+
+  ```python
+  @qml.qnode(qml.device("lightning.qubit", wires=1))
+  def f(y):
+      qml.RX(y, wires=0)
+      return qml.expval(qml.PauliZ(0))
+
+  @catalyst.qjit
+  def g(x):
+      return catalyst.grad(lambda y: f(jnp.cos(y)) ** 2)(x)
+  ```
+
+  ```pycon
+  >>> g(0.4)
+  0.3751720385067584
+  ```
+
 <h3>Internal changes</h3>
 
 * Remove the `MemMemCpyOptPass` in llvm O2 (applied for Enzyme), this reduces bugs when 

frontend/test/pytest/test_gradient.py

@@ -34,6 +34,7 @@
     pure_callback,
     qjit,
     value_and_grad,
+    vmap,
 )
 
 # pylint: disable=too-many-lines
@@ -1444,7 +1445,25 @@ def interpreted(x):
     assert np.allclose(compiled(inp), interpreted(inp))
 
 
-@pytest.mark.xfail(reason="Needs PR #332")
+@pytest.mark.parametrize("inp", [(1.0), (2.0), (3.0), (4.0)])
+def test_preprocessing_outside_qnode(inp, backend):
+    """Test the preprocessing outside qnode."""
+
+    @qml.qnode(qml.device(backend, wires=1))
+    def f(y):
+        qml.RX(y, wires=0)
+        return qml.expval(qml.PauliZ(0))
+
+    @qjit
+    def g(x):
+        return grad(lambda y: f(jnp.cos(y)) ** 2)(x)
+
+    def h(x):
+        return jax.grad(lambda y: f(jnp.cos(y)) ** 2)(x)
+
+    assert np.allclose(g(inp), h(inp))
+
+
 def test_gradient_slice(backend):
     """Test the differentation when the qnode generates memref with non identity layout."""
     n_wires = 5
@@ -1484,6 +1503,60 @@ def my_model(data, weights, bias):
     jax_res = jax.jacobian(my_model, argnums=1)(data, params["weights"], params["bias"])
     assert np.allclose(cat_res, jax_res)
 
+@pytest.mark.xfail(reason="Vmap yields wrong results when differentiated")
+def test_vmap_worflow_derivation(backend):
+    """Check the gradient of a vmap workflow"""
+    n_wires = 5
+    data = jnp.sin(jnp.mgrid[-2:2:0.2].reshape(n_wires, -1)) ** 3
+
+    targets = jnp.array([-0.2, 0.4, 0.35, 0.2], dtype=jax.numpy.float64)
+
+    dev = qml.device(backend, wires=n_wires)
+
+    @qml.qnode(dev, diff_method="adjoint")
+    def circuit(data, weights):
+        """Quantum circuit ansatz"""
+
+        @for_loop(0, n_wires, 1)
+        def data_embedding(i):
+            qml.RY(data[i], wires=i)
+
+        data_embedding()
+
+        @for_loop(0, n_wires, 1)
+        def ansatz(i):
+            qml.RX(weights[i, 0], wires=i)
+            qml.RY(weights[i, 1], wires=i)
+            qml.RX(weights[i, 2], wires=i)
+            qml.CNOT(wires=[i, (i + 1) % n_wires])
+
+        ansatz()
+
+        return qml.expval(qml.sum(*[qml.PauliZ(i) for i in range(n_wires)]))
+
+    circuit = vmap(circuit, in_axes=(1, None))
+
+    def my_model(data, weights, bias):
+        return circuit(data, weights) + bias
+
+    def loss_fn(params, data, targets):
+        predictions = my_model(data, params["weights"], params["bias"])
+        loss = jnp.sum((targets - predictions) ** 2 / len(data))
+        return loss
+
+    weights = jnp.ones([n_wires, 3])
+    bias = jnp.array(0.0, dtype=jax.numpy.float64)
+    params = {"weights": weights, "bias": bias}
+
+    results_enzyme = qjit(grad(loss_fn))(params, data, targets)
+    results_jax = jax.grad(loss_fn)(params, data, targets)
+
+    data_enzyme, pytree_enzyme = tree_flatten(results_enzyme)
+    data_jax, pytree_fd = tree_flatten(results_jax)
+
+    assert pytree_enzyme == pytree_fd
+    assert jnp.allclose(data_enzyme[0], data_jax[0])
+    assert jnp.allclose(data_enzyme[1], data_jax[1])
 
 @pytest.mark.parametrize(
     "gate,state", ((qml.BasisState, np.array([1])), (qml.StatePrep, np.array([0, 1])))

frontend/test/pytest/test_vmap.py

@@ -18,8 +18,11 @@
 import jax.numpy as jnp
 import pennylane as qml
 import pytest
+from jax.tree_util import tree_flatten
 
-from catalyst import qjit, vmap
+from catalyst import for_loop, grad, qjit, vmap
+
+# pylint: disable=too-many-public-methods
 
 
 class TestVectorizeMap:

mlir/lib/Gradient/Transforms/GradMethods/HybridGradient.cpp

@@ -197,7 +197,8 @@ static FailureOr<func::FuncOp> cloneCallee(PatternRewriter &rewriter, Operation
                 funcOp.walk([&](func::CallOp callOp) {
                     if (callOp.getCallee() == qnode.getName()) {
                         PatternRewriter::InsertionGuard insertionGuard(rewriter);
-                        rewriter.setInsertionPointToStart(&funcOp.getFunctionBody().front());
+                        // TODO: optimize the placement of the param count call (e.g. loop hoisting)
+                        rewriter.setInsertionPoint(callOp);
                         Value paramCount =
                             rewriter
                                 .create<func::CallOp>(loc, paramCountFn, callOp.getArgOperands())