Fix compilation of Jacobian with for loops and vmap

frontend/catalyst/jit.py

@@ -384,6 +384,7 @@ def compute_jvp(self, primals, tangents):
         # Optimization: Do not compute Jacobians for arguments which do not participate in
         #               differentiation.
         argnums = []
+        tangents, _ = tree_flatten(tangents)
         for idx, tangent in enumerate(tangents):
             if not isinstance(tangent, jax.custom_derivatives.SymbolicZero):
                 argnums.append(idx)

frontend/test/pytest/test_vmap.py

@@ -18,8 +18,9 @@
 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
 
 
 class TestVectorizeMap:
@@ -752,3 +753,57 @@ def circuit(x):
             match="Invalid batch size; it must be a non-zero integer, but got 0.",
         ):
             qjit(workflow)(x)
+
+    def test_vmap_worflow_derivation(self, 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])
+
+        dev = qml.device(backend, wires=n_wires)
+
+        @qml.qnode(dev)
+        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 = qml.qjit(vmap(circuit, in_axes=(1, None)), autograph=False)
+
+        def my_model(data, weights, bias):
+            return circuit(data, weights) + bias
+
+        @qml.qjit
+        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)
+        params = {"weights": weights, "bias": bias}
+
+        results_enzyme = qml.qjit(grad(loss_fn))(params, data, targets)
+        results_fd = jax.grad(loss_fn)(params, data, targets)
+
+        data_enzyme, pytree_enzyme = tree_flatten(results_enzyme)
+        data_fd, pytree_fd = tree_flatten(results_fd)
+
+        assert pytree_enzyme == pytree_fd
+        assert jnp.allclose(data_enzyme[0], data_fd[0])
+        assert jnp.allclose(data_enzyme[1], data_fd[1])

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

@@ -183,7 +183,7 @@ FailureOr<func::FuncOp> HybridGradientLowering::cloneCallee(PatternRewriter &rew
                 funcOp.walk([&](func::CallOp callOp) {
                     if (callOp.getCallee() == qnode.getName()) {
                         PatternRewriter::InsertionGuard insertionGuard(rewriter);
-                        rewriter.setInsertionPointToStart(&funcOp.getFunctionBody().front());
+                        rewriter.setInsertionPoint(callOp);
                         Value paramCount =
                             rewriter
                                 .create<func::CallOp>(loc, paramCountFn, callOp.getArgOperands())