Merge pull request #13602 from mattjj:tweak-array-notebook

PiperOrigin-RevId: 494480882

CHANGELOG.md

@@ -18,8 +18,7 @@ Remember to align the itemized text with the first line of an item within a list
     breaking change to the `pjit` API.  The [jax.Array migration
     guide](https://jax.readthedocs.io/en/latest/jax_array_migration.html) can
     help you migrate your codebase to `jax.Array`. You can also look at the
-    [Parallelism with
-    JAX](https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html)
+    [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html)
     tutorial to understand the new concepts.
   * `PartitionSpec` and `Mesh` are now out of experimental. The new API endpoints
     are `jax.sharding.PartitionSpec` and `jax.sharding.Mesh`.

docs/conf.py

@@ -210,7 +210,7 @@
     # TODO(jakevdp): enable execution on the following if possible:
     'jax-101/*',
     'notebooks/xmap_tutorial.*',
-    'notebooks/Parallelism_with_JAX.*',
+    'notebooks/Distributed_arrays_and_automatic_parallelization.*',
 ]
 
 # -- Options for HTMLHelp output ---------------------------------------------

docs/index.rst

@@ -7,7 +7,7 @@ parallelize, Just-In-Time compile to GPU/TPU, and more.
 
 .. note::
    JAX 0.4.0 introduces new parallelism APIs, including breaking changes to :func:`jax.experimental.pjit` and a new unified ``jax.Array`` type.
-   Please see `Parallelism with JAX <https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html>`_ tutorial and the :ref:`jax-array-migration`
+   Please see `Distributed arrays and automatic parallelization <https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html>`_ tutorial and the :ref:`jax-array-migration`
    guide for more information.
 
 .. toctree::
@@ -52,7 +52,7 @@ parallelize, Just-In-Time compile to GPU/TPU, and more.
 
    notebooks/autodiff_cookbook
    multi_process
-   notebooks/Parallelism_with_JAX
+   notebooks/Distributed_arrays_and_automatic_parallelization
    notebooks/vmapped_log_probs
    notebooks/neural_network_with_tfds_data
    notebooks/Custom_derivative_rules_for_Python_code

docs/jax_array_migration.md

@@ -18,7 +18,7 @@ the unified jax.Array
 After the migration is complete `jax.Array` will be the only type of array in
 JAX.
 
-This doc explains how to migrate existing codebases to `jax.Array`. For more information on using `jax.Array` and JAX parallelism APIs, see the [Parallelism with JAX](https://jax.readthedocs.io/en/latest/notebooks/Parallelism_with_JAX.html) tutorial.
+This doc explains how to migrate existing codebases to `jax.Array`. For more information on using `jax.Array` and JAX parallelism APIs, see the [Distributed arrays and automatic parallelization](https://jax.readthedocs.io/en/latest/notebooks/Distributed_arrays_and_automatic_parallelization.html) tutorial.
 
 
 ### How to enable jax.Array?