update_data_loading_on_cpu

Author: selamw1

docs/data_loaders_on_cpu_with_jax.ipynb

@@ -6,7 +6,7 @@
     "id": "PUFGZggH49zp"
    },
    "source": [
-    "## Introduction to Data Loaders on CPU with JAX"
+    "# Introduction to Data Loaders on CPU with JAX"
    ]
   },
   {
@@ -15,6 +15,8 @@
     "id": "3ia4PKEV5Dr8"
    },
    "source": [
+    "[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jax-ml/jax-ai-stack/blob/main/docs/data_loaders_on_cpu_with_jax.ipynb)\n",
+    "\n",
     "This tutorial explores different data loading strategies for using **JAX** on a single [**CPU**](https://jax.readthedocs.io/en/latest/glossary.html#term-CPU). While JAX doesn't include a built-in data loader, it seamlessly integrates with popular data loading libraries, including:\n",
     "*   [**PyTorch DataLoader**](https://github.com/pytorch/data)\n",
     "*   [**TensorFlow Datasets (TFDS)**](https://github.com/tensorflow/datasets)\n",
@@ -262,7 +264,7 @@
     "id": "Hsionp5IYsQ9"
    },
    "source": [
-    "## Loading Data with PyTorch DataLoader\n",
+    "### Loading Data with PyTorch DataLoader\n",
     "\n",
     "This section shows how to load the MNIST dataset using PyTorch's DataLoader, convert the data to NumPy arrays, and apply transformations to flatten and cast images."
    ]
@@ -614,7 +616,7 @@
     "id": "Nm45ZTo6yrf5"
    },
    "source": [
-    "## Loading Data with TensorFlow Datasets (TFDS)\n",
+    "### Loading Data with TensorFlow Datasets (TFDS)\n",
     "\n",
     "This section demonstrates how to load the MNIST dataset using TFDS, fetch the full dataset for evaluation, and define a training generator for batch processing. GPU usage is explicitly disabled for TensorFlow."
    ]
@@ -816,7 +818,7 @@
     "id": "-ryVkrAITS9Z"
    },
    "source": [
-    "## Loading Data with Grain\n",
+    "### Loading Data with Grain\n",
     "\n",
     "This section demonstrates how to load MNIST data using Grain, a data-loading library. You'll define a custom dataset class for Grain and set up a Grain DataLoader for efficient training."
    ]
@@ -1082,7 +1084,7 @@
     "id": "oixvOI816qUn"
    },
    "source": [
-    "## Loading Data with Hugging Face\n",
+    "### Loading Data with Hugging Face\n",
     "\n",
     "This section demonstrates loading MNIST data using the Hugging Face `datasets` library. You'll format the dataset for JAX compatibility, prepare flattened images and one-hot-encoded labels, and define a training generator."
    ]

docs/data_loaders_on_cpu_with_jax.md

@@ -5,18 +5,20 @@ jupytext:
     extension: .md
     format_name: myst
     format_version: 0.13
-    jupytext_version: 1.16.4
+    jupytext_version: 1.15.2
 kernelspec:
   display_name: Python 3
   name: python3
 ---
 
 +++ {"id": "PUFGZggH49zp"}
 
-## Introduction to Data Loaders on CPU with JAX
+# Introduction to Data Loaders on CPU with JAX
 
 +++ {"id": "3ia4PKEV5Dr8"}
 
+[![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/jax-ml/jax-ai-stack/blob/main/docs/data_loaders_on_cpu_with_jax.ipynb)
+
 This tutorial explores different data loading strategies for using **JAX** on a single [**CPU**](https://jax.readthedocs.io/en/latest/glossary.html#term-CPU). While JAX doesn't include a built-in data loader, it seamlessly integrates with popular data loading libraries, including:
 *   [**PyTorch DataLoader**](https://github.com/pytorch/data)
 *   [**TensorFlow Datasets (TFDS)**](https://github.com/tensorflow/datasets)
@@ -188,7 +190,7 @@ def train_model(num_epochs, params, training_generator, data_loader_type='stream
 
 +++ {"id": "Hsionp5IYsQ9"}
 
-## Loading Data with PyTorch DataLoader
+### Loading Data with PyTorch DataLoader
 
 This section shows how to load the MNIST dataset using PyTorch's DataLoader, convert the data to NumPy arrays, and apply transformations to flatten and cast images.
 
@@ -318,7 +320,7 @@ train_model(num_epochs, params, pytorch_training_generator(mnist_dataset), data_
 
 +++ {"id": "Nm45ZTo6yrf5"}
 
-## Loading Data with TensorFlow Datasets (TFDS)
+### Loading Data with TensorFlow Datasets (TFDS)
 
 This section demonstrates how to load the MNIST dataset using TFDS, fetch the full dataset for evaluation, and define a training generator for batch processing. GPU usage is explicitly disabled for TensorFlow.
 
@@ -413,7 +415,7 @@ train_model(num_epochs, params, training_generator)
 
 +++ {"id": "-ryVkrAITS9Z"}
 
-## Loading Data with Grain
+### Loading Data with Grain
 
 This section demonstrates how to load MNIST data using Grain, a data-loading library. You'll define a custom dataset class for Grain and set up a Grain DataLoader for efficient training.
 
@@ -547,7 +549,7 @@ train_model(num_epochs, params, pygrain_training_generator)
 
 +++ {"id": "oixvOI816qUn"}
 
-## Loading Data with Hugging Face
+### Loading Data with Hugging Face
 
 This section demonstrates loading MNIST data using the Hugging Face `datasets` library. You'll format the dataset for JAX compatibility, prepare flattened images and one-hot-encoded labels, and define a training generator.