Jax 0.8 seems to break the manipulation example notebook

Hi, thanks for the wonderful repository! I was trying to get the manipulation.ipynb example working and after training the LeapCubeReorient policy, the fingers would always disappear. After many hours of debugging, it seems that downgrading to jax[cuda]==0.7.2 solved the issue. Reproducing the environment here for anyone else that faces this issue:

```
uv venv --python 3.12.2
source .venv/bin/activate
uv pip install jax[cuda]==0.7.2
uv pip install mujoco mujoco_mjx brax
uv pip install mediapy
uv pip install playground
```