Fix links in DTensor Overview guide and ML Tutorial.

The `.ipynb` suffix triggers a download of the notebook instead of opening it on the Tensorflow.org site. Also fixed a broken link in the ML Tutorial and remove extraneous whitespace.

PiperOrigin-RevId: 446235732

Author: srujun

site/en/guide/dtensor_overview.ipynb

@@ -74,7 +74,7 @@
         "\n",
         "DTensor provides a global programming model that allows developers to compose applications that operate on Tensors globally while managing the distribution across devices internally. DTensor distributes the program and tensors according to the sharding directives through a procedure called *[Single program, multiple data (SPMD)](https://en.wikipedia.org/wiki/SPMD) expansion*.\n",
         "\n",
-        "By decoupling the application from sharding directives, DTensor enables running the same application on a single device, multiple devices, or even multiple clients, while preserving its global semantics. \n",
+        "By decoupling the application from sharding directives, DTensor enables running the same application on a single device, multiple devices, or even multiple clients, while preserving its global semantics.\n",
         "\n",
         "This guide introduces DTensor concepts for distributed computing, and how DTensor integrates with TensorFlow. To see a demo of using DTensor in model training, see [Distributed training with DTensor](https://www.tensorflow.org/tutorials/distribute/dtensor_ml_tutorial) tutorial."
       ]
@@ -130,7 +130,7 @@
         "  tf.config.set_logical_device_configuration(phy_devices[0], [\n",
         "        tf.config.LogicalDeviceConfiguration(),\n",
         "    ] * ncpu)\n",
-        "  \n",
+        "\n",
         "configure_virtual_cpus(6)\n",
         "DEVICES = [f'CPU:{i}' for i in range(6)]\n",
         "\n",
@@ -147,8 +147,8 @@
         "\n",
         "DTensor introduces two concepts: `dtensor.Mesh` and `dtensor.Layout`. They are abstractions to model the sharding of tensors across topologically related devices.\n",
         "\n",
-        "- `Mesh` defines the device list for computation. \n",
-        "- `Layout` defines how to shard the Tensor dimension on a `Mesh`. "
+        "- `Mesh` defines the device list for computation.\n",
+        "- `Layout` defines how to shard the Tensor dimension on a `Mesh`."
       ]
     },
     {
@@ -223,7 +223,7 @@
         "\n",
         "**`Layout`** specifies how a tensor is distributed, or sharded, on a `Mesh`. In DTensor, the placement specification is on a per-axis bases. An axis of a `Layout` can be either `sharded` or `unsharded` (replicated) along a mesh dimension.\n",
         "\n",
-        "Note: In order to avoid confusions between `Mesh` and `Layout`, the term *dimension* is always associated with `Mesh`, and the term *axis* with `Tensor` and `Layout` in this guide. \n",
+        "Note: In order to avoid confusions between `Mesh` and `Layout`, the term *dimension* is always associated with `Mesh`, and the term *axis* with `Tensor` and `Layout` in this guide.\n",
         "\n",
         "The rank of a `Layout` and the number of dimensions of a `Mesh` do not need to match. The `unsharded` axes of a `Layout` do not need to be associated to a mesh dimension, and `unsharded` mesh dimensions do not need to be associated with a `layout` axis.\n",
         "\n",
@@ -298,9 +298,9 @@
         "\n",
         "DTensor supports both single-client and multi-client applications. The colab Python kernel is an example of a single client DTensor application, where there is a single Python process.\n",
         "\n",
-        "In a multi-client DTensor application, multiple Python processes collectively perform as a coherent application. The Cartisian grid of a `Mesh` in a multi-client DTensor application can span across devices regardless of whether they are attached locally to the current client or attached remotely to another client. The set of all devices used by a `Mesh` are called the *global device list*. \n",
+        "In a multi-client DTensor application, multiple Python processes collectively perform as a coherent application. The Cartisian grid of a `Mesh` in a multi-client DTensor application can span across devices regardless of whether they are attached locally to the current client or attached remotely to another client. The set of all devices used by a `Mesh` are called the *global device list*.\n",
         "\n",
-        "The creation of a `Mesh` in a multi-client DTensor application is a collective operation where the *global device list* is identicial for all of the participating clients, and the creation of the `Mesh` serves as a global barrier. \n",
+        "The creation of a `Mesh` in a multi-client DTensor application is a collective operation where the *global device list* is identicial for all of the participating clients, and the creation of the `Mesh` serves as a global barrier.\n",
         "\n",
         "During `Mesh` creation, each client provides its *local device list* together with the expected *global device list*. DTensor validates that both lists are consistent. Please refer to the API documentation for `dtensor.create_mesh` and `dtensor.create_distributed_mesh`\n",
         " for more information on multi-client mesh creation and the *global device list*.\n",
@@ -331,17 +331,17 @@
       "source": [
         "def dtensor_from_array(arr, layout, shape=None, dtype=None):\n",
         "  \"\"\"Convert a DTensor from something that looks like an array or Tensor.\n",
-        "  \n",
+        "\n",
         "  This function is convenient for quick doodling DTensors from a known,\n",
         "  unsharded data object in a single-client environment. This is not the\n",
-        "  most efficient way of creating a DTensor, but it will do for this \n",
+        "  most efficient way of creating a DTensor, but it will do for this\n",
         "  tutorial.\n",
         "  \"\"\"\n",
         "  if shape is not None or dtype is not None:\n",
         "    arr = tf.constant(arr, shape=shape, dtype=dtype)\n",
-        "  \n",
+        "\n",
         "  # replicate the input to the mesh\n",
-        "  a = dtensor.copy_to_mesh(arr, \n",
+        "  a = dtensor.copy_to_mesh(arr,\n",
         "          layout=dtensor.Layout.replicated(layout.mesh, rank=layout.rank))\n",
         "  # shard the copy to the desirable layout\n",
         "  return dtensor.relayout(a, layout=layout)"
@@ -356,11 +356,11 @@
         "### Anatomy of a DTensor\n",
         "\n",
         "A DTensor is a `tf.Tensor` object, but augumented with the `Layout` annotation that defines its sharding behavior. A DTensor consists of the following:\n",
-        "  \n",
+        "\n",
         "  - Global tensor meta-data, including the global shape and dtype of the tensor.\n",
         "  - A `Layout`, which defines the `Mesh` the `Tensor` belongs to, and how the `Tensor` is sharded onto the `Mesh`.\n",
-        "  - A list of **component tensors**, one item per local device in the `Mesh`. \n",
-        "  \n",
+        "  - A list of **component tensors**, one item per local device in the `Mesh`.\n",
+        "\n",
         "With `dtensor_from_array`, you can create your first DTensor, `my_first_dtensor`, and examine its contents."
       ]
     },
@@ -450,7 +450,8 @@
       "source": [
         "The inverse operation of `dtensor.unpack` is `dtensor.pack`. Component tensors can be packed back into a DTensor.\n",
         "\n",
-        "The components must have the same rank and dtype, which will be the rank and dtype of the returned DTensor. However there is no strict requirement on the device placement of component tensors as inputs of `dtensor.unpack`: the function will automatically copy the component tensors to their respective corresponding devices. \n"
+        "The components must have the same rank and dtype, which will be the rank and dtype of the returned DTensor. However there is no strict requirement on the device placement of component tensors as inputs of `dtensor.unpack`: the function will automatically copy the component tensors to their respective corresponding devices.\n",
+        "\n"
       ]
     },
     {
@@ -504,7 +505,7 @@
         "Create a 3x2 rank-2 DTensor, sharding its first axis along the `'x'` mesh dimension, and its second axis along the `'y'` mesh dimension.\n",
         "\n",
         "- Because the tensor shape equals to the mesh dimension along all of the sharded axes, each device receives a single element of the DTensor.\n",
-        "- The rank of the component tensor is always the same as the rank of the global shape. DTensor adopts this convention as a simple way to preserve information for locating the relation between a component tensor and the global DTensor. "
+        "- The rank of the component tensor is always the same as the rank of the global shape. DTensor adopts this convention as a simple way to preserve information for locating the relation between a component tensor and the global DTensor."
       ]
     },
     {
@@ -567,10 +568,10 @@
         "DTensor allows a `Layout` to be a hybrid, sharded along some axes, but replicated along others.\n",
         "\n",
         "For example, you can shard the same 3x2 rank-2 DTensor in the following way:\n",
-        "  \n",
+        "\n",
         "  - 1st axis sharded along the `'x'` mesh dimension.\n",
         "  - 2nd axis replicated along the `'y'` mesh dimension.\n",
-        "  \n",
+        "\n",
         "To achieve this sharding scheme, you just need to replace the sharding spec of the 2nd axis from `'y'` to `dtensor.UNSHARDED`, to indicate your intention of replicating along the 2nd axis. The layout object will look like `Layout(['x', dtensor.UNSHARDED], mesh)`."
       ]
     },
@@ -610,7 +611,7 @@
       "source": [
         "#### Tensor.numpy() and sharded DTensor\n",
         "\n",
-        "Be aware that calling the `.numpy()` method on a sharded DTensor raises an error. The rationale for erroring is to protect against unintended gathering of data from multiple computing devices to the host CPU device backing the returned numpy array. "
+        "Be aware that calling the `.numpy()` method on a sharded DTensor raises an error. The rationale for erroring is to protect against unintended gathering of data from multiple computing devices to the host CPU device backing the returned numpy array."
       ]
     },
     {
@@ -624,12 +625,12 @@
         "print(fully_replicated_dtensor.numpy())\n",
         "\n",
         "try:\n",
-        "  fully_sharded_dtensor.numpy() \n",
+        "  fully_sharded_dtensor.numpy()\n",
         "except tf.errors.UnimplementedError:\n",
         "  print(\"got an error as expected for fully_sharded_dtensor\")\n",
         "\n",
         "try:\n",
-        "  hybrid_sharded_dtensor.numpy() \n",
+        "  hybrid_sharded_dtensor.numpy()\n",
         "except tf.errors.UnimplementedError:\n",
         "  print(\"got an error as expected for hybrid_sharded_dtensor\")"
       ]
@@ -650,11 +651,11 @@
         "  - Rewriting TensorFlow Ops on the global DTensor with equivalent TensorFlow Ops on the componenent tensors, inserting collective and communication Ops when necessary\n",
         "  - Lowering backend neutral TensorFlow Ops to backend specific TensorFlow Ops.\n",
         "\n",
-        "The final result is that **DTensor is a drop-in replacement for Tensor**. \n",
+        "The final result is that **DTensor is a drop-in replacement for Tensor**.\n",
         "\n",
         "Note: DTensor is still an experimental API which means you will be exploring and pushing the boundaries and limits of the DTensor programming model.\n",
         "\n",
-        "There are 2 ways of triggering DTensor execution: \n",
+        "There are 2 ways of triggering DTensor execution:\n",
         "  - DTensor as operands of a Python function, e.g. `tf.matmul(a, b)` will run through DTensor if `a`, `b`, or both are DTensors.\n",
         "  - Requesting the result of a Python function to be a DTensor, e.g. `dtensor.call_with_layout(tf.ones, layout, shape=(3, 2))` will run through DTensor because we requested the output of tf.ones to be sharded according to a `layout`."
       ]
@@ -678,10 +679,10 @@
       "source": [
         "#### Fully replicated input and output\n",
         "\n",
-        "In this case, the DTensors are fully replicated. On each of the devices of the `Mesh`, \n",
+        "In this case, the DTensors are fully replicated. On each of the devices of the `Mesh`,\n",
         "  - the component tensor for operand `a` is `[[1, 2, 3], [4, 5, 6]]` (2x3)\n",
         "  - the component tensor for operand `b` is `[[6, 5], [4, 3], [2, 1]]` (3x2)\n",
-        "  - the computation consists of a single `MatMul` of `(2x3, 3x2) -> 2x2`, \n",
+        "  - the computation consists of a single `MatMul` of `(2x3, 3x2) -> 2x2`,\n",
         "  - the component tensor for result `c` is `[[20, 14], [56,41]]` (2x2)\n",
         "\n",
         "Total number of floating point mul operations is `6 device * 4 result * 3 mul = 72`."
@@ -792,7 +793,7 @@
         "\n",
         "What about Python functions that do not take operands, but returns a Tensor result that can be sharded? Examples of such functions are\n",
         "\n",
-        "  - `tf.ones`, `tf.zeros`, `tf.random.stateless_normal`, \n",
+        "  - `tf.ones`, `tf.zeros`, `tf.random.stateless_normal`,\n",
         "\n",
         "For these Python functions, DTensor provides `dtensor.call_with_layout` which eagelry executes a Python function with DTensor, and ensures that the returned Tensor is a DTensor with the requested `Layout`."
       ]
@@ -885,9 +886,9 @@
       "outputs": [],
       "source": [
         "ones = dtensor.call_with_layout(\n",
-        "    tf.function(tf.random.stateless_normal), \n",
-        "    dtensor.Layout(['x', 'y'], mesh), \n",
-        "    shape=(6, 4), \n",
+        "    tf.function(tf.random.stateless_normal),\n",
+        "    dtensor.Layout(['x', 'y'], mesh),\n",
+        "    shape=(6, 4),\n",
         "    seed=(1, 1))\n",
         "print(ones)"
       ]
@@ -910,8 +911,8 @@
       "outputs": [],
       "source": [
         "ones = dtensor.call_with_layout(\n",
-        "    tf.function(tf.ones), \n",
-        "    dtensor.Layout(['x', 'y'], mesh), \n",
+        "    tf.function(tf.ones),\n",
+        "    dtensor.Layout(['x', 'y'], mesh),\n",
         "    shape=(6, 4))\n",
         "print(ones)"
       ]
@@ -925,7 +926,7 @@
         "### From `tf.Variable` to `dtensor.DVariable`\n",
         "\n",
         "In Tensorflow, `tf.Variable` is the holder for a mutable `Tensor` value.\n",
-        "With DTensor, the corresponding variable semantics is provided by `dtensor.DVariable`. \n",
+        "With DTensor, the corresponding variable semantics is provided by `dtensor.DVariable`.\n",
         "\n",
         "The reason a new type `DVariable` was introduced for DTensor variable is because DVariables have an additional requirement that the layout cannot change from its initial value."
       ]
@@ -945,7 +946,7 @@
         "    initial_value=dtensor.call_with_layout(\n",
         "        tf.function(tf.random.stateless_normal),\n",
         "        layout=layout,\n",
-        "        shape=tf.TensorShape([64, 32]), \n",
+        "        shape=tf.TensorShape([64, 32]),\n",
         "        seed=[1, 1],\n",
         "        dtype=tf.float32))\n",
         "\n",