Jax Issues Fix?

I noticed a lot of issues submitted related to jax function deprecations. Just a suggestion here that changing the dm-haiku dependency in ai-models-graphcast pyproject.toml from 0.0.10 to 0.0.13 seems to solve all of them.