Update to jax 0.4.19

[jsbrittain]

Provide compatbility with jax > 0.4.14
Tag jax and jaxlib to version 0.4.19

Resolves #10

[dfm]

Thanks for this!! I've had this on my to do list for ages, but never seem to get around to it. Much appreciated!

Could you add a sentence at the end of the first section of the README (below the bit about "January 2021" and "use at your own risk") to note this update?

[dfm]

Looks like we'll also need to update the Python version (maybe 3.10) for the CI workflow:

extending-jax/.github/workflows/tests.yml

Line 26 in 1cb4c39

python-version: 3.8

[jsbrittain]

Looks like we'll also need to update the Python version (maybe 3.10) for the CI workflow:

extending-jax/.github/workflows/tests.yml

Line 26 in 1cb4c39

python-version: 3.8

I've also updated the Dockerfile to ubuntu 22.04 for the same reason.

[dfm]

Thanks again!!