Closed dime10 closed 2 weeks ago
Thanks @dime10! How urgent would you say this is to fix?
Hey @trbromley, I would likely label this a P0/High Priority, as it is blocking an external researcher
Might this be a lightning bug? Because it works with default.qubit
dev = qml.device("default.qubit", wires=1)
@jax.jit
@qml.qnode(dev, interface="jax")
def circuit1(theta):
qml.exp(qml.X(0), -1j * jnp.cos(theta))
return qml.expval(qml.Z(0))
circuit1(jnp.pi / 3.)
Array(0.54030231, dtype=float64)
Might this be a lightning bug? Because it works with default.qubit
dev = qml.device("default.qubit", wires=1) @jax.jit @qml.qnode(dev, interface="jax") def circuit1(theta): qml.exp(qml.X(0), -1j * jnp.cos(theta)) return qml.expval(qml.Z(0)) circuit1(jnp.pi / 3.)
Array(0.54030231, dtype=float64)
In that case the difference could be in how the gate is executed. Probably the matrix
method is traceable but not the decomposition
method, I believe lightning favours the latter whereas default favours the former.
@Qottmann @dime10 I can confirm that the issue only occurs when combining jax.jit
with workflows that trigger Exp.decomposition
.
The underlying issue is afaik the Exp._recursive_decomposition
and Exp.has_generator
methods.
In both methods, Python boolean expressions are being made with qml.math.real(coeff)
. This works fine without JIT, but when using JIT, qml.math.real(coeff)
will be a dynamic variable and Python does not know how to convert this into a bool.
A couple of notes:
You can replace the boolean expression with
jnp.logical_and(self.base.is_hermitian, jnp.logical_not(qml.math.real(self.coeff)))
This will work fine (for JAX), but not for other frameworks. The output of this expression will now be a JAX dynamic variable with boolean type.
However, even with this change, Exp._recursive_decomposition
will still fail, because you can't use Python if
conditions with a JAX dynamic tracer:
Finally, even if Exp._recursive_decomposition
and Exp.has_generator
were re-written to either work with JAX or remove the control flow that depends on dynamic variables, there's a chance this will still fail because JAX does not support recursive functions of dynamic variables.
Expected behavior
Originally reported in https://github.com/PennyLaneAI/catalyst/issues/923
Actual behavior
The following circuit raises an error:
Additional information
No response
Source code
No response
Tracebacks
No response
System information
Existing GitHub issues