Issues with JAX CUDA Installation Instructions and Deprecated Sharding API in Code

Hello DeepMind team,

I was excited to explore your repository and replicate the setup for some experiments. However, I encountered two issues during installation and code review that should be addressed to improve compatibility and maintainability. I'm reporting them here with details for reproducibility.

#### 1. Outdated JAX CUDA Installation Command for Newer CUDA Versions
The README.md suggests the following command for installing JAX with CUDA support (under the "Installation" section):

```
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

This command fails with newer CUDA versions (e.g., CUDA 12.6 or later, which are common in recent setups). The error typically relates to incompatible wheel availability or build issues.

**Proposed Fix:** Update the command to use the more general CUDA installer, which automatically detects the CUDA version:

```
pip install jax[cuda] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
```

This worked seamlessly on my system with CUDA 12.6. For users with specific versions, you could add a note to adjust based on their setup, but the general form should be the default recommendation. Reference: JAX installation docs (https://jax.readthedocs.io/en/latest/installation.html#pip-installation-gpu-cuda).

#### 2. Use of Deprecated `jax.sharding.PositionalSharding` API
In the codebase, `jax.sharding.PositionalSharding` is used. This API has been deprecated in recent JAX versions in favor of `jax.sharding.NamedSharding` combined with `jax.sharding.PartitionSpec` for better flexibility and future-proofing.

**Proposed Fix:** Replace instances of `PositionalSharding` with `PartitionSpec`. This ensures compatibility with JAX's evolving sharding API. If needed, I can provide a pull request with the changes once I pinpoint all occurrences (I spotted at least one in the training-related utils).

These fixes would make the repo more accessible for users with modern hardware and software stacks. Let me know if you need more details or logs from my setup (Ubuntu 20.04.6, Python 3.12.10, JAX 0.8.0, CUDA 12.6).

Thanks for the great work on amortized planning, looking forward to seeing updates!

Best,
Zeev

Start with the Installation section of README.md, then search the training-related utils for jax.sharding.PositionalSharding. Confirm the affected occurrences and the supported replacement approach before changing the CUDA command and sharding usage; done means the documented installation path and identified deprecated API uses are addressed.

Written by the indexing model from the issue text.