Welcome to the very first lesson of JAX in Action: Building an Image Classifier! Congratulations on reaching this final course in our comprehensive JAX learning path — you've come incredibly far and should be proud of your dedication and progress.

Let's quickly recap the amazing journey that brought us here. In our first course, we mastered JAX fundamentals, learning how NumPy-like operations work in JAX's functional programming paradigm and discovering the power of automatic differentiation with jax.grad. Our second course took us deeper into advanced JAX concepts, where we explored efficient batching with jax.vmap, flexible data structures with PyTrees, and the performance benefits of JIT compilation. The third course equipped us with essential deep learning tools, teaching us to build neural networks from scratch using pure JAX, implement training loops with gradient descent, and solve classic problems like XOR. Finally, our fourth course introduced us to the elegant world of Flax and Optax, where we learned to construct modular neural networks, leverage sophisticated optimizers like Adam, and create production-ready training pipelines.

Now, in this final course, we're ready to tackle a real-world challenge: building a complete image classification system from the ground up. Throughout the course, we'll work with a real dataset, implement convolutional neural networks, and apply everything we've learned to solve practical computer vision problems. Today's first lesson focuses on establishing our project foundation and mastering the critical first step of any machine learning project: data loading and preprocessing.

The MNIST dataset serves as our gateway into image classification: think of it as the "Hello, World!" of computer vision. MNIST contains 70,000 handwritten digit images (0-9), each measuring 28×28 pixels in grayscale. Despite its simplicity, MNIST teaches us fundamental concepts that scale to complex real-world datasets: pixel normalization, shape manipulation, and batch processing.

We'll use TensorFlow Datasets (tensorflow_datasets or tfds) to access MNIST, even though we're building with JAX. This might seem counterintuitive, but tfds provides an excellent data loading ecosystem with built-in preprocessing, caching, and performance optimizations. Many JAX practitioners use this combination because TensorFlow's data pipeline integrates seamlessly with JAX arrays through NumPy compatibility.

The beauty of tfds lies in its standardized interface. Whether we're working with MNIST, CIFAR-10, or ImageNet, the loading patterns remain consistent. This consistency will serve us well when we expand to more complex datasets in later lessons. The library handles downloading, caching, and version management automatically, so we can focus on the machine learning rather than data infrastructure.

Let's begin constructing our data loading module, which will serve as the backbone of our entire pipeline. We'll start by an empty src/__init__.py file to define the Python module and implementing our data core preprocessing logic in src/data_loader.py:

The preprocess_image function handles the critical data transformations that prepare raw MNIST data for neural network consumption. First, we convert images from unsigned 8-bit integers (0-255) to 32-bit floats, then normalize by dividing by 255.0 to scale pixel values into the [0,1] range. This normalization is essential because neural networks train more effectively with inputs in standardized ranges.

The reshaping operation adds a channel dimension, transforming each image from shape (28, 28) to (28, 28, 1). This seemingly minor change is crucial for convolutional neural networks, which expect images in height-width-channels format. Even though MNIST images are grayscale (single channel), we explicitly represent this dimension for consistency with color image formats that may come later.

Now we'll build the full data loading pipeline that handles batching, shuffling, and performance optimization:

This function demonstrates production-grade data pipeline construction. The tfds.load call retrieves both training and test splits while capturing dataset metadata in ds_info. We use as_supervised=False to maintain dictionary format for our examples, giving us explicit control over the data structure.

The training data pipeline applies several performance optimizations: map with parallel processing applies our preprocessing function efficiently, cache stores preprocessed data in memory to avoid redundant computation, shuffle randomizes example order to prevent learning biases (notice that we specify a seed), batch groups examples for efficient GPU processing, and prefetch overlaps data loading with model training to minimize idle time.

Let's finish our data loader:

The test data pipeline follows a similar pattern but omits shuffling — we want consistent evaluation results across runs. The crucial final step converts TensorFlow datasets to NumPy iterators using tfds.as_numpy, bridging the gap between TensorFlow's data loading and JAX's computation engine.

Finally, we can test our data pipeline using the src/main.py file. We'll examine what our data pipeline produces and verify everything works correctly:

This main module imports the data loading functionality and initializes the pipeline with a batch size of 32. This batch size balances memory usage with training efficiency — small enough to fit comfortably in memory while large enough to provide stable gradient estimates during training. Then, we extract a sample image and related label for inspection.

When we execute this code, we'll see comprehensive information about our data pipeline:

This output confirms our pipeline works perfectly! The dataset contains 60,000 training examples and 10,000 test examples. Each image is properly shaped as (28, 28, 1), and our preprocessing successfully normalized pixel values to the [0,1] range. The batch dimensions show we're getting 32 images per batch, exactly as configured, with the correct data types for JAX consumption.

Excellent work! You've successfully established the foundation of a professional machine learning project with clean modular architecture and robust data preprocessing. We've created a reusable data loading system that efficiently handles the MNIST dataset while following industry best practices for performance and maintainability.

The pipeline we've built demonstrates several critical ML engineering principles: separation of concerns through modular design, efficient data processing with TensorFlow's optimized operations, and seamless integration between different frameworks. In our next lesson, we'll build upon this foundation by implementing our first convolutional neural network using Flax, bringing us one step closer to building a complete image classification system.