[Frontend] Support for transforms

[erick-xanadu]

Context: Allow PL transforms to be traced in Catalyst

Description of the Change: These changes are implemented on top of the tape-transform PR in Catalyst and the development version of PL. These changes apply transforms that are stored on the device into the tape, potentially generating many tapes. Then each tape is traced individually. At last, the classical co-transform is also traced.

TODO:

• These changes do not include yet the "one function per tape" invariant.
• Nested transforms

Benefits: Catalyst support for PL transforms.

[sc-41344]

[rmoyard]

Great PR, very exciting 💯 Some suggestions but nothing blocking. If you pin to the latest PennyLane and the tests pass. I am happy to approve!

[erick-xanadu]

@rmoyard since this depends on the main branch of PL and we currently have a undiscovered bug when integrating with the main branch of PL, I'll leave this for now.

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Comparison is base (7c00ec4) 99.63% compared to head (597f2a4) 99.63%.

Additional details and impacted files

@@           Coverage Diff           @@
##             main     #280   +/-   ##
=======================================
  Coverage   99.63%   99.63%           
=======================================
  Files          44       44           
  Lines        7666     7734   +68     
  Branches      455      467   +12     
=======================================
+ Hits         7638     7706   +68     
  Misses         14       14           
  Partials       14       14           

Files	Coverage Δ
frontend/catalyst/jax_tracer.py	99.68% <100.00%> (+0.08%)	⬆️
frontend/catalyst/pennylane_extensions.py	99.64% <100.00%> (+<0.01%)	⬆️

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

[rmoyard]

Thank you @erick-xanadu great job 💯 It would be nice to gather the possible improvements (relax the return type, support composable transforms) with issues and TODO in the code. Otherwise it looks all good!

[josh146]

Nice work @erick-xanadu! Just a few flyby questions/comments

• Don't forget to add the changelog entry!

• Don't forget to also update this section of the docs: https://docs.pennylane.ai/projects/catalyst/en/latest/dev/sharp_bits.html#compatibility-with-pennylane-transforms

• I'm curious if everything still works if Autograph is turned on?

[erick-xanadu]

I'm curious if everything still works if Autograph is turned on?

Transforms are restricted to straight-line programs so I don't think autograph would make much sense here.

[dime10]

Great work @erick-xanadu!! 💯

[dime10]

Great 💯