How does a neural network turn an image into predictions? A forward pass applies each layer in order: multiply activations by weights, add biases, apply a nonlinear activation, and produce final output activations. In JAX, this computation is a stateless function such as predict(params, image), with model parameters passed explicitly rather than stored on an object. section 2.4.2 section 2.4
Classifying a flattened image
A flattened 28 × 28 grayscale image provides 784 input values. Passing it to predict(params, random_flattened_image) produces an output with shape (10,), representing the network's ten target-class activations.







