Support arbitrary generating terms in pulse Hamiltonians with stoch_pulse_grad

[dwierichs]

Context:
Thus far, the gradient transform stoch_pulse_grad only supports Pauli words as generating terms in the pulse Hamiltonian.

Description of the Change:
This PR generalizes the support to arbitrary generating terms, using the existing generalized shift rule framework.
The inserted operations are no longer PauliRot gates but Exp operators, because for non-Pauli word generators, the decomposition into Pauli rotation gates is not always desired.
Decomposing Exp operations with Pauli words in the exponent yields back the PauliRot operations used by stoch_pulse_grad until now.

Benefits:
Support many more pulse Hamiltonians, in particular e.g. neutral atom Hamiltonians.

Possible Drawbacks:
Breaking change: For Pauli word-generated terms, the inserted operations are now Exp operations instead of PauliRot gates.

Open TODOs

• Add example of non-Pauli word generator to docstring.
• Add JIT-compatibility test with non-Pauli word generator.

[sc-34236]

[codecov]

Codecov Report

Merging #4132 (fc4adc5) into master (6ab9dda) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4132   +/-   ##
=======================================
  Coverage   99.77%   99.77%           
=======================================
  Files         340      340           
  Lines       30595    30607   +12     
=======================================
+ Hits        30525    30537   +12     
  Misses         70       70           

Impacted Files	Coverage Δ
pennylane/gradients/pulse_gradient.py	100.00% <100.00%> (ø)
pennylane/ops/op_math/symbolicop.py	100.00% <100.00%> (ø)

[Qottmann]

Very nice upgrade @dwierichs 💪

Much appreciate the developer comments 🤩

There is a slight performance regression on the first execution as ParametrizedHamiltonian gets jit compiled inside the ODE solver (because the ode fun is jit-compiled), most likely due to the extra steps from qml.exp, so dont think there is much we can do. Could short-cut the execution by checking whether it is a pauli word, but that might make the code significantly uglier.

Execution time seems unchanged 👍

[dwierichs]

Could short-cut the execution by checking whether it is a pauli word, but that might make the code significantly uglier.

I tried this and it seems to be fine to me, actually 😇 Let me know what you think :)

edit: The definition of qml.pauli.is_pauli_word includes rescaled Pauli words. Therefore, we will need a function that extracts a rescaling parameter from a Pauli word (which is allowed to be a Hamiltonian, SProd, Prod, Tensor or atomic Pauli operator). I think it's nice to add this, but I'd do it in a separate PR, in case we want the simplification Exp->PauliRot for the inserted operations.

[Qottmann]

Very neat upgrade with not many changes in code and tests, really nice work @dwierichs 👌

With the upcoming short-cutting for Pauli-words looks good to me 🙂

[dwierichs]

@Qottmann while fixing the bug in #4156 I found a simple way to extract the prefactor of a Pauli word that I didn't see before. I used this to implement the PauliRot shortcut for Pauli words in this PR :)
This includes single-term Hamiltonians, SProd instances and (tensor products of) simple Pauli operators.

[rmoyard]

Couple of questions and suggestion, otherwise it looks good 💯

[Qottmann]

Thanks for fixing the short-cut for pauli words 👌 Jit times look much better compared to current master for those cases. Execution time is the same 👍

[rmoyard]

Thanks a lot! 👍