Merge pull request #4 from trallard/trallard/patch-pr-110

Fix headings hierarchy

Author: selamw1

docs/data_loaders_for_multi_device_setups_with_jax.ipynb

@@ -57,7 +57,7 @@
     "id": "TsFdlkSZKp9S"
    },
    "source": [
-    "### Checking TPU Availability for JAX"
+    "## Checking TPU Availability for JAX"
    ]
   },
   {
@@ -99,7 +99,7 @@
     "id": "qyJ_WTghDnIc"
    },
    "source": [
-    "### Setting Hyperparameters and Initializing Parameters\n",
+    "## Setting Hyperparameters and Initializing Parameters\n",
     "\n",
     "You'll define hyperparameters for your model and data loading, including layer sizes, learning rate, batch size, and the data directory. You'll also initialize the weights and biases for a fully-connected neural network."
    ]
@@ -141,7 +141,7 @@
     "id": "rHLdqeI7D2WZ"
    },
    "source": [
-    "### Model Prediction with Auto-Batching\n",
+    "## Model Prediction with Auto-Batching\n",
     "\n",
     "In this section, you'll define the `predict` function for your neural network. This function computes the output of the network for a single input image.\n",
     "\n",
@@ -210,7 +210,7 @@
     "id": "rLqfeORsERek"
    },
    "source": [
-    "### Utility and Loss Functions\n",
+    "## Utility and Loss Functions\n",
     "\n",
     "You'll now define utility functions for:\n",
     "- One-hot encoding: Converts class indices to binary vectors.\n",

docs/data_loaders_for_multi_device_setups_with_jax.md

@@ -44,7 +44,7 @@ from jax.sharding import Mesh, PartitionSpec, NamedSharding
 
 +++ {"id": "TsFdlkSZKp9S"}
 
-### Checking TPU Availability for JAX
+## Checking TPU Availability for JAX
 
 ```{code-cell}
 ---
@@ -58,7 +58,7 @@ jax.devices()
 
 +++ {"id": "qyJ_WTghDnIc"}
 
-### Setting Hyperparameters and Initializing Parameters
+## Setting Hyperparameters and Initializing Parameters
 
 You'll define hyperparameters for your model and data loading, including layer sizes, learning rate, batch size, and the data directory. You'll also initialize the weights and biases for a fully-connected neural network.
 
@@ -90,7 +90,7 @@ params = init_network_params(layer_sizes, random.PRNGKey(0))
 
 +++ {"id": "rHLdqeI7D2WZ"}
 
-### Model Prediction with Auto-Batching
+## Model Prediction with Auto-Batching
 
 In this section, you'll define the `predict` function for your neural network. This function computes the output of the network for a single input image.
 
@@ -139,7 +139,7 @@ sharding_spec = PartitionSpec('device')
 
 +++ {"id": "rLqfeORsERek"}
 
-### Utility and Loss Functions
+## Utility and Loss Functions
 
 You'll now define utility functions for:
 - One-hot encoding: Converts class indices to binary vectors.