Drawing the circuit with catalyst.measure fails

Issue description

Description of the issue - Circuit drawing fails when using mid circuit measurment.

• Expected behavior: (What you expect to happen)
  Draw the circuit

• Actual behavior: (What actually happens)
  Error

• 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.34.0
  Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
  Home-page: https://github.com/PennyLaneAI/pennylane
  Author:
  Author-email:
  License: Apache License 2.0
  Location: /envs/catalyst/lib/python3.10/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: Linux-6.2.0-37-generic-x86_64-with-glibc2.35
  Python version: 3.10.13
  Numpy version: 1.26.3
  Scipy version: 1.12.0
  Installed devices:

• lightning.qubit (PennyLane-Lightning-0.34.0)
• default.gaussian (PennyLane-0.34.0)
• default.mixed (PennyLane-0.34.0)
• default.qubit (PennyLane-0.34.0)
• default.qubit.autograd (PennyLane-0.34.0)
• default.qubit.jax (PennyLane-0.34.0)
• default.qubit.legacy (PennyLane-0.34.0)
• default.qubit.tf (PennyLane-0.34.0)
• default.qubit.torch (PennyLane-0.34.0)
• default.qutrit (PennyLane-0.34.0)
• null.qubit (PennyLane-0.34.0)

Source code and tracebacks

@qjit
@qml.qnode(qml.device("lightning.qubit", wires=2,shots=100))
def circuit(theta):
    qml.Hadamard(wires=0)
    qml.RX(theta, wires=1)
    qml.CNOT(wires=[0,1])
    m=measure(wires=0)
    return qml.expval(qml.PauliZ(wires=1)),m

print(qml.draw(circuit)(1.2))

Error:
catalyst.utils.exceptions.CompileError: catalyst.measure can only be used from within @qjit.

Start by running the provided circuit with qml.draw(circuit), then inspect how @qjit and catalyst.measure interact at the drawing boundary. The issue is done when the reproduced circuit draws successfully without the “can only be used from within @qjit” CompileError.