Merge pull request #20713 from jakevdp:parallelism-doc

PiperOrigin-RevId: 623877877

Author: jax authors

docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb

@@ -17,11 +17,7 @@
    "source": [
     "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/google/jax/blob/main/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb) [![Open in Kaggle](https://kaggle.com/static/images/open-in-kaggle.svg)](https://kaggle.com/kernels/welcome?src=https://github.com/google/jax/blob/main/docs/notebooks/Distributed_arrays_and_automatic_parallelization.ipynb)\n",
     "\n",
-    "This tutorial discusses parallelism via `jax.Array`, the unified array object model available in JAX v0.4.1 and newer.\n",
-    "\n",
-    "Refer to the [`jax.Array migration`](https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration) guide to learn how to migrate the existing JAX pre-v0.4.1 codebases to `jax.Array`.\n",
-    "\n",
-    "**Note:** The features required by `jax.Array` are not supported by the Colab TPU runtime at this time, but are available on Google Cloud TPU and Kaggle TPU VMs."
+    "This tutorial discusses parallelism via `jax.Array`, the unified array object model available in JAX v0.4.1 and newer."
    ]
   },
   {
@@ -2369,6 +2365,7 @@
   }
  ],
  "metadata": {
+  "accelerator": "TPU",
   "colab": {
    "provenance": [],
    "toc_visible": true

docs/notebooks/Distributed_arrays_and_automatic_parallelization.md

@@ -21,10 +21,6 @@ kernelspec:
 
 This tutorial discusses parallelism via `jax.Array`, the unified array object model available in JAX v0.4.1 and newer.
 
-Refer to the [`jax.Array migration`](https://jax.readthedocs.io/en/latest/jax_array_migration.html#jax-array-migration) guide to learn how to migrate the existing JAX pre-v0.4.1 codebases to `jax.Array`.
-
-**Note:** The features required by `jax.Array` are not supported by the Colab TPU runtime at this time, but are available on Google Cloud TPU and Kaggle TPU VMs.
-
 ```{code-cell}
 :id: FNxScTfq3vGF