Add changes to support Jax 0.4

[vcharraut]

Description

Add changes to support jax>=0.4.0 and jaxlib>=0.4.0

Fixes #359

Note : There might be some typing imprecisions, I think some Any type could be remove since jax.Array is supposed to be general, but I wasn't so sure.

To add in the next major release

Type of change

Please delete options that are not relevant.

• New feature (non-breaking change which adds functionality)

Screenshots

Please attach before and after screenshots of the change if applicable.

Checklist:

• I have run the pre-commit checks with pre-commit run --all-files (see CONTRIBUTING.md instructions to set it up)
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[pseudo-rnd-thoughts]

I think the direction we should take with this is to wait for v0.29 but we should keep this open until then
The reason is that python 3.7 still has 3 months until end of life so if just wait until to closer to the time

[pseudo-rnd-thoughts]

@charraut There has been a change of plan and the next release with be a major release allowing us to include this PR
Could resolve the conflicts then we should be good to merge

[vcharraut]

@charraut There has been a change of plan and the next release with be a major release allowing us to include this PR Could resolve the conflicts then we should be good to merge

Done! Let me know if you need any specific changes

[pseudo-rnd-thoughts]

Done! Let me know if you need any specific changes

Amazing, thanks