Batched csr in ParametrizedEvolution with jax version bump

[Qottmann]

Following #4079 which hot-fixes #4078, we can now fix it for good with batched CSR. This update is necessary anyway since "vanilla" CSR format will be deprecated soon (see jax-ml/jax#15844 (comment))

[codecov]

Codecov Report

Merging #4126 (3086b56) into master (d95ae12) will not change coverage.
The diff coverage is 100.00%.

@@           Coverage Diff           @@
##           master    #4126   +/-   ##
=======================================
  Coverage   99.78%   99.78%           
=======================================
  Files         340      340           
  Lines       30523    30523           
=======================================
  Hits        30456    30456           
  Misses         67       67           

Impacted Files	Coverage Δ
pennylane/pulse/parametrized_hamiltonian_pytree.py	100.00% <100.00%> (ø)

[dwierichs]

Nice change, I like how minimal it is while unlocking quite a bit of performance in the long run :)
Just wondering whether we

1. may break anything with the new doc/requirements settings, and
2. should add a test that makes use of a particular BCSR feature that CSR would fail.
   For the second, I think it is not really necessary, but it might be nice.

[rmoyard]

Thanks @Qottmann it looks good to me 👍

[dwierichs]

Looks good to me, thanks @Qottmann ! 🎉

[Qottmann]

should add a test that makes use of a particular BCSR feature that CSR would fail.
For the second, I think it is not really necessary, but it might be nice.

Yes good point! Testing now differentiating the unitary matrix with jit, no jit, dense, and sparse :)