Make compilation transforms support jax.jit (reapply changes)

[antalszava]

Context:
Reapplies the changes from #1894.

[codecov]

Codecov Report

Merging #2397 (8c20b30) into master (910e95b) will decrease coverage by 0.00%.
The diff coverage is 100.00%.

@@            Coverage Diff             @@
##           master    #2397      +/-   ##
==========================================
- Coverage   99.43%   99.42%   -0.01%     
==========================================
  Files         245      245              
  Lines       18896    18884      -12     
==========================================
- Hits        18789    18776      -13     
- Misses        107      108       +1     

Impacted Files	Coverage Δ
.../transforms/decompositions/single_qubit_unitary.py	100.00% <100.00%> (ø)
...ane/transforms/decompositions/two_qubit_unitary.py	100.00% <100.00%> (ø)
...nnylane/transforms/optimization/merge_rotations.py	97.61% <100.00%> (-2.39%)	⬇️
...lane/transforms/optimization/optimization_utils.py	100.00% <100.00%> (ø)
...ane/transforms/optimization/single_qubit_fusion.py	100.00% <100.00%> (ø)

---

Continue to review full report at Codecov.

Legend - Click here to learn more
Δ = absolute <relative> (impact), ø = not affected, ? = missing data
Powered by Codecov. Last update 910e95b...8c20b30. Read the comment docs.

[albi3ro]

Glad you caught this. Looks fine to go for me, assuming the original PR had a thorough review.

[antalszava]

Thank you! All thanks for catching it goes to @glassnotes 🎉