A vector is a rank-1 tensor: a multidimensional array with one axis. Its size counts elements, while its rank counts dimensions. A classification model’s output for one input is a vector of class probabilities, with one element per class. section 2.5
Vectors also appear inside larger computations. A batch of predictions adds another dimension, producing a matrix rather than a vector; JAX’s vmap can apply a single-input prediction function across that batch. In parallel workloads, vector elements can be partitioned across devices, and dot products combine corresponding products from the shards. section 2.5 section 8.1.3
From one prediction to a batch
A single prediction is a vector of class scores. Applying the vmapped function to 32 inputs produces an array with shape (32, 10): one 10-element prediction vector for each input. section 2.5
Listing 2.9 Using the vmapped function
batched_preds = batched_predict(params, random_flattened_images) #1 print(batched_preds.shape) >>> (32, 10) #2
Inspect the vmapped prediction call and its resulting batch shape. section 2.5
