Autograph fails when used with Catalyst's JAX VJP integration

[isaacdevlugt]

Issue description

The following circuit

import jax
import pennylane as qml

num_qubits = 4
H = qml.X(0) @ qml.X(1) @ qml.Y(2) @ qml.Z(3)
dev = qml.device("lightning.qubit", wires=num_qubits)


@qml.qjit(autograph=True)
@qml.qnode(dev)
def cost(weights):
    for i, x in enumerate(weights):
        qml.RY(x, wires=[i])

    return qml.expval(H)


weights = jax.numpy.ones(num_qubits)
results = jax.grad(cost)(weights)
print(results)

produces the following warning

WARNING:root:AutoGraph could not transform <function outer_factory.<locals>.inner_factory.<locals>.ag__cost at 0x291887d80> and will run it as-is.
Please report this to the TensorFlow team. When filing the bug, set the verbosity to 10 (on Linux, `export AUTOGRAPH_VERBOSITY=10`) and attach the full output.
Cause: closure mismatch, requested ('ag__',), but source function had ()

indicating that AutoGraph is not converting the function.

• Expected behavior: (What you expect to happen)

Calls to jax.grad or similar should work, but you'd have to use catalyst.grad instead to successfully qjit the code.

• Actual behavior: (What actually happens)

See above warning message

• Reproduces how often: (What percentage of the time does it reproduce?)

100%

• System information: (post the output of import pennylane as qml; qml.about())

Name: PennyLane
Version: 0.35.1
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/isaac/.virtualenvs/pennylane-catalyst/lib/python3.11/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning

Platform info:           macOS-14.4.1-arm64-arm-64bit
Python version:          3.11.8
Numpy version:           1.26.4
Scipy version:           1.13.0
Installed devices:
- default.clifford (PennyLane-0.35.1)
- default.gaussian (PennyLane-0.35.1)
- default.mixed (PennyLane-0.35.1)
- default.qubit (PennyLane-0.35.1)
- default.qubit.autograd (PennyLane-0.35.1)
- default.qubit.jax (PennyLane-0.35.1)
- default.qubit.legacy (PennyLane-0.35.1)
- default.qubit.tf (PennyLane-0.35.1)
- default.qubit.torch (PennyLane-0.35.1)
- default.qutrit (PennyLane-0.35.1)
- null.qubit (PennyLane-0.35.1)
- nvidia.custatevec (PennyLane-Catalyst-0.5.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.5.0)
- softwareq.qpp (PennyLane-Catalyst-0.5.0)
- lightning.qubit (PennyLane_Lightning-0.35.1)

Source code and tracebacks

See above