fix: adapt to new jax API

[GaetanLepage]

Tests are currently broken when ran with the latest version of JAX installed.
This patch ensures the code compatibility with the current jax API.

[borisdayma]

Thanks, that's all you needed to change?
I'm scared the training code will most likely be broken with latest version as well.

I pinned the versions to use here:

dalle-mini/setup.cfg

Line 32 in f0be4de

jax==0.3.25

[GaetanLepage]

Thanks, that's all you needed to change?

Well at least, it is enough for all of the tests to pass.
For the context, I was updating dalle-mini to the latest version on nixpkgs.