[BUG] parameterizing input coefficients of `qml.exp` fails with qjit

Experienced error when trying to parameterize qml.exp with qjit. Example code below.

Expected behavior:

import pennylane as qml
import jax.numpy as jnp
from catalyst import qjit

@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(theta):
    qml.exp(qml.X(0), coeff=-1j*jnp.cos(theta))
    return qml.expval(qml.Z(0))

circuit(jnp.pi/3)

Output:

>>> 0.5403023058681395

Actual behavior: with qjit

import pennylane as qml
import jax.numpy as jnp
from catalyst import qjit

@qjit
@qml.qnode(qml.device('lightning.qubit', wires=1))
def circuit(theta):
    qml.exp(qml.X(0), coeff=-1j*jnp.cos(theta))
    return qml.expval(qml.Z(0))

circuit(jnp.pi/3)

output:

---------------------------------------------------------------------------
TracerBoolConversionError                 Traceback (most recent call last)
[<ipython-input-139-a7dc14c8fdfe>](https://localhost:8080/#) in <cell line: 11>()
      9     return qml.expval(qml.Z(0))
     10 
---> 11 circuit(jnp.pi/3)

36 frames
[/usr/local/lib/python3.10/dist-packages/jax/_src/core.py](https://localhost:8080/#) in error(self, arg)
   1508   if fun is bool:
   1509     def error(self, arg):
-> 1510       raise TracerBoolConversionError(arg)
   1511   elif fun in (hex, oct, operator.index):
   1512     def error(self, arg):

TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[]..
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

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

Name: PennyLane
Version: 0.37.0
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: /usr/local/lib/python3.10/dist-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning

Platform info:           Linux-6.1.85+-x86_64-with-glibc2.35
Python version:          3.10.12
Numpy version:           1.25.2
Scipy version:           1.11.4
Installed devices:
- default.clifford (PennyLane-0.37.0)
- default.gaussian (PennyLane-0.37.0)
- default.mixed (PennyLane-0.37.0)
- default.qubit (PennyLane-0.37.0)
- default.qubit.autograd (PennyLane-0.37.0)
- default.qubit.jax (PennyLane-0.37.0)
- default.qubit.legacy (PennyLane-0.37.0)
- default.qubit.tf (PennyLane-0.37.0)
- default.qubit.torch (PennyLane-0.37.0)
- default.qutrit (PennyLane-0.37.0)
- default.qutrit.mixed (PennyLane-0.37.0)
- default.tensor (PennyLane-0.37.0)
- null.qubit (PennyLane-0.37.0)
- nvidia.custatevec (PennyLane-Catalyst-0.7.0)
- nvidia.cutensornet (PennyLane-Catalyst-0.7.0)
- oqc.cloud (PennyLane-Catalyst-0.7.0)
- softwareq.qpp (PennyLane-Catalyst-0.7.0)
- lightning.qubit (PennyLane_Lightning-0.37.0)

Additional information: replacing qml.exp(qml.X(0), coeff=-1j*jnp.cos(theta)) with qml.evolve(qml.X(0), coeff = jnp.cos(theta)) causes the same error TracerBoolConversionError, but replacing it with qml.RX(2*jnp.cos(theta),0) doesn't cause the error to occur.

PennyLaneAI / catalyst

[BUG] parameterizing input coefficients of `qml.exp` fails with qjit #923