Need a compact, familiar task for testing a classifier? MNIST supplies handwritten-digit images for multiclass classification experiments. section 2.1
A typical JAX workflow chooses MNIST, loads it into batches, trains a model, and uses the resulting parameters for predictions. section 2.2
Figure 2.1 High-level structure of a JAX project, with data loading, training, and deployment to production
Trace MNIST from data loading through training and deployment in a high-level JAX project structure. section 2.2




