Enable dense keyword in ParametrizedEvolution

[Qottmann]

Hot fix for #4078 (see also google/jax#15844)

Enables a keyword dense to be passed to ParametrizedEvolution that enforces dense evolution, which currently automatically is enabled for more than 2 wires

Edit2: will have to wait for jax 0.4.9 support to enable BCSR, so this is a hotfix for the moment (see discussion)

[github-actions]

Hello. You may have forgotten to update the changelog!
Please edit doc/releases/changelog-dev.md with:

• A one-to-two sentence description of the change. You may include a small working example for new features.
• A link back to this PR.
• Your name (or GitHub username) in the contributors section.

[codecov]

Codecov Report

Merging #4079 (02c4e1e) into master (755b241) will increase coverage by 0.00%.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4079   +/-   ##
=======================================
  Coverage   99.71%   99.71%           
=======================================
  Files         345      345           
  Lines       30737    30738    +1     
=======================================
+ Hits        30650    30651    +1     
  Misses         87       87           

Impacted Files	Coverage Δ
pennylane/pulse/parametrized_evolution.py	95.94% <100.00%> (+0.05%)	⬆️

[Qottmann]

Implemented the fix from google/jax#15844 and #4078

In case that solves it (and doesnt cause any unexpected behavior / failing tests), would we still want the possibility to enforce dense matrices?

[PhilipVinc]

Can we also have a test checking that specifying dense=True works?

[PhilipVinc]

cause any unexpected behavior / failing tests), would we still want the possibility to enforce dense matrices?

I'd say that's useful. Can be deprecated in the future if we don't want this to be an option anymore.

[Qottmann]

This is blocked because we are stuck at 0.4.3 jax version. The main problem has been fixed on the jax side in google/jax#14833 but we would also need to drop python 3.8 and that may take a bit longer 🥲 CC @rmoyard

[mlxd]

Just a fly-by comment: once this PR is ready to go can we enable the GPU CI tests by adding the "gpu" label to this PR?
To avoid overheads, I'd consider this only when the work is ready to merge.

[Qottmann]

Thanks, would have forgotten that! Are GPU tests running all jax tests with gpu enabled? 🤩 I always wanted to test the pulse stuff on GPU but never had the occasion

[rmoyard]

The GPU tests for the moment only have Torch.

[mlxd]

Yep, as @rmoyard said only torch is supported now. We should ensure the runner installs the JAX CUDA version for this to be supported.

[dwierichs]

Thanks for the quick fix @Qottmann 💪
Three small comments :)

[dwierichs]

Thanks @Qottmann !

[Qottmann]

Merging this and waiting with the change from csr to batched csr format when we have jax 0.4.9 supported in PL