Bump jax and jaxlib versions used in tests.
Bump jax and jaxlib versions used in tests.

PiperOrigin-RevId: 342229266
Change-Id: Ibf3d78ed5477ccde5ae0d3f0739eb29d0d54eb9f

deepmind/dm-haiku,deepmind/dm-haiku