[Frontend] JAX tracer supporting user-defined tape transformation

[grwlf]

In this PR we aim at making it possible to apply user-defined transformations to a tape containing Catalyst quantum operators. The main difficulty is related to the fact that in Catalyst we only trace programs rather than execute them so the operations on tape: (1) are filled with abstract parameters, (2) include control flow operations containing nested tapes.

The approach we are working on would essentially split the tracing process into classical and quantum parts. In order to do so, we store JAX tracing contexts and re-enter them as needed. We described the idea and the associated risks in the internal document.

Tasks:

• Prototyping is done
• Core tracing framework is implemented
• Catalyst extensions are ported to the new tracing framework
  • Control flow ops (For-loops, WhileLoops, Conditionals, Adjoint)
  • CustomOps
  • QubitUnitary
  • Mid-circuit measurement
  • Measurements (Sample, ExpVal, Counts, Probs, State)
  • Observables (Compbasis, Paulis, Hadamard, Hermitian,Tensor,Hamiltonian)
  • Gradient primitives (Grad, JVP, VJP)
• All kinds of execution pipelines are supported
  • Both qjit and qnode
  • Just qjit (no quantum parts)
  • Just qnode (interpretation)
• Optimizations are ported
  • Qubit insert/extract state preservation (quantum register tracing optimisation)
  • Differentiation-specific procedures (JVP/VJP/Grad) are reviewed and ported
• All tests pass locally
• Latest version of Catalyst is merged into the PR
  • Obvious merge conflicts are solves
  • Pytests are fixed
  • Littests are fixed
• Code is clean
  • Temporary and unused source files are removed
  • JAX internals are collected in a separate utility file
  • Tracing context classes are unified
  • Pylint is happy
  • Docstrings are ported or created
  • CodeFactor is happy
  • Codecov is happy
  • Remaining FIXMEs are fixed

[sc-41345]

[codecov]

Codecov Report

Patch coverage: 99.62% and project coverage change: +0.04% 🎉

Comparison is base (8106beb) 99.47% compared to head (02e1361) 99.51%.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #240      +/-   ##
==========================================
+ Coverage   99.47%   99.51%   +0.04%     
==========================================
  Files          41       40       -1     
  Lines        7192     7196       +4     
  Branches      386      410      +24     
==========================================
+ Hits         7154     7161       +7     
+ Misses         20       17       -3     
  Partials       18       18              

Files Changed	Coverage Δ
frontend/catalyst/pennylane_extensions.py	99.58% <99.34%> (+0.43%)	⬆️
frontend/catalyst/jax_tracer.py	99.55% <99.50%> (+0.27%)	⬆️
frontend/catalyst/compilation_pipelines.py	100.00% <100.00%> (ø)
frontend/catalyst/jax_primitives.py	97.08% <100.00%> (+0.03%)	⬆️
frontend/catalyst/utils/contexts.py	100.00% <100.00%> (ø)
frontend/catalyst/utils/jax_extras.py	100.00% <100.00%> (ø)

... and 1 file with indirect coverage changes

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[dime10]

Looks good to me 💯