Open trbromley opened 1 month ago
@trbromley do you have more details on this?
There is also an error when using qml.qjit and lightning.qubit, which doesn't expect a JAX tracer for the seed argument.
This PR might be relevant, which adds a seed kwarg to qjit: https://github.com/PennyLaneAI/catalyst/pull/936
@trbromley do you have more details on this?
There is also an error when using qml.qjit and lightning.qubit, which doesn't expect a JAX tracer for the seed argument.
When I run this:
import jax
import pennylane as qml
def circuit(key, param):
dev = qml.device("lightning.qubit", wires=2, shots=10, seed=key)
@qml.qnode(dev, interface="jax")
def my_circuit():
qml.RX(param, wires=0)
qml.CNOT(wires=[0, 1])
return qml.sample(qml.PauliZ(0))
return my_circuit()
key = jax.random.PRNGKey(1967)
I get this error:
>>> qml.qjit(circuit)(key, 0.5)
TypeError: SeedSequence expects int or sequence of ints for entropy not Traced<ShapedArray(uint32[2])>with<DynamicJaxprTrace(level=1/0)>
This PR might be relevant, which adds a seed kwarg to qjit: PennyLaneAI/catalyst#936
Oh interesting, so the seed should live with qjit
rather than the device? :thinking:
Oh interesting, so the seed should live with qjit rather than the device? 🤔
The PR was motivated by a technical problem: we have tests in the catalyst code base which are flaky, but no qjit-compatible way of specifying a seed. So passing it through the qjit
decorator is a quick way of getting this support in there, but we can revisit from a user POV if we need to
I get this error:
Ah yes I guess this is to be expected, since JAX seeds are not supported by MLIR/LLVM
Expected behavior
It is possible to JAX JIT a function where
qml.device
hasseed
set to be a tracer for the PRNG key.Actual behavior
An error is raised in many cases. I'm not exactly sure when, but it appears to be if there are other tracers active - like with the
param
variable below.There is also an error when using
qml.qjit
andlightning.qubit
, which doesn't expect a JAX tracer for theseed
argument.Additional information
Although there is a draft PR to fix this problem, it has gone stale and I decided it would be more useful to track the problem in this issue.
Source code
Tracebacks
System information
Existing GitHub issues