Bump version to 0.2.

Adding xarray_jax vmap.

Automated Code Change

[graphcast] Migrate to the new `jax.pmap`.

Ensure correct sharding of rng during rollout.

Replace unicode escaped characters in ipynb files

Fix feature utils for update to scipy 1.17

[graphcast] Prepare for `jax_pmap_shmap_merge=True`.

[pmap] Remove `jax_pmap_shmap_merge` configuration.

Update contact information with privacy policy.

Update test for `jax_pmap_shmap_merge=True`.

Replace `jax.stages.OutInfo` with `jax.ShapeDtypeStruct`

Clean up math in `get_rotation_matrices_to_local_coordinates`.

Allowing any leading dims in einsum.

Explicitly inject numpy and scipy.spatial.transform in some model utils.