This issue reflexes the fixme comment in the quantum tracing pipeline for the conditional primitive.
The core of the problem is: we re-trace Jaxpr program in order to unify the result values (link) but at the same time, for quantum programs, we use two-staged classical-quantum tracing. Before quantum tracing is complete, the Jaxpr IR might not be valid from the Jax point of view. Some tracers might not have their source equations yet, e.g. this might happen with programs involving qml.measure. Because of this, we only perform the unification during the quantum tracing (link), which might be too late.
In order to fix this problem, we might need to either (1) guarantee that the IR is always valid after the classical tracing is complete which might be hard or (2) Implement the unification "manually" in a way that would allow dangling Jaxpr variables.
The second approach seems be more realistic in the short-term.
This issue reflexes the fixme comment in the quantum tracing pipeline for the conditional primitive.
The core of the problem is: we re-trace Jaxpr program in order to unify the result values (link) but at the same time, for quantum programs, we use two-staged classical-quantum tracing. Before quantum tracing is complete, the Jaxpr IR might not be valid from the Jax point of view. Some tracers might not have their source equations yet, e.g. this might happen with programs involving
qml.measure. Because of this, we only perform the unification during the quantum tracing (link), which might be too late.In order to fix this problem, we might need to either (1) guarantee that the IR is always valid after the classical tracing is complete which might be hard or (2) Implement the unification "manually" in a way that would allow dangling Jaxpr variables.
The second approach seems be more realistic in the short-term.
See also: