[jax update, main] Update jax to 0.5.3

[paul0403]

Context: This is the main branch for updating jax to 0.5.3.

We will do this version-by-version. Each version should come in as a PR/commit. This allows us to run CI frequently at each version.

Current highest successful version: 0.5.3 (without lapack)

[sc-88694]