Open erick-xanadu opened 1 month ago
What is the bug?
Note, could we instead of changing jax.interpreters.partial_eval.DynamicJaxprTracer and adding a hash, can't we change pennylane wire utilities to find whether the wire is jax.interpreters.partial_eval.DynamicJaxprTracer and compute the id as its hash instead of modifying jax itself?
Yes I think that's a good idea! The capture module will probably take care of this.
What is the bug?
The bug is that we are changing the global state and we shouldn't :sweat_smile: Just changing the state when we shouldn't. The concrete error is #894 but that is just a symptom.
There are several places where we modify JAX during compilation.
Just to list some:
Patchers... (see jax_extras, see jax_transient_config)
And we also have a global context to see whether or not we are running or jax via the EvaluationContext.
With callbacks, this now changes the assumption that if we are tracing, we will never go back to the python environment. We should have a function that is able to save the configuration before we trace, change however we want it, reset it during callbacks, and reset it back to what we need once we exit the callback scope.
Note, could we instead of changing
jax.interpreters.partial_eval.DynamicJaxprTracer
and adding a hash, can't we change pennylane wire utilities to find whether the wire isjax.interpreters.partial_eval.DynamicJaxprTracer
and compute theid
as itshash
instead of modifying jax itself?