Skip to content

← JAX Vectorization & Control Flow step 3 of 25

Easy Primitives

vmap with None broadcasting

Jump to the task

Share weights across a batch

A model usually applies the same weights to many examples. With vmap, mark which argument supplies the examples and which stays shared. For arguments (W, x_batch), in_axes=(None, 0) means:

  • Use the complete W for every example.
  • Map over axis 0 of x_batch, passing one row at a time to the function being transformed.

This describes the result as if you applied the function separately to each example. vmap transforms the array operations; it does not promise to run a Python loop or call your Python function once per row. None does not add a dimension or request physical copies of W.

Worked example

import jax
import jax.numpy as jnp

W = jnp.array([[1.0, 0.0],
               [0.0, 1.0],
               [1.0, 1.0]])          # shape (3, 2)
x_batch = jnp.array([[1.0, 2.0],
                     [3.0, 4.0]])    # shape (2, 2)

apply_batch = jax.vmap(lambda w, x: w @ x, in_axes=(None, 0))
result = apply_batch(W, x_batch)
# [[1., 2., 3.],
#  [3., 4., 7.]]

For the first example, the three rows of W select the first feature, select the second feature, and add both features. This gives [1, 2, 3]. The second example gives [3, 4, 7]. Each output has shape (3,); the default out_axes=0 stacks the two outputs into shape (2, 3).

Problem

Implement affine_per_example(W, x_batch) using jax.vmap with in_axes=(None, 0) to apply W @ x to each example.

  • W has shape (d_out, d_in).
  • x_batch has shape (N, d_in).
  • Return shape (N, d_out), with out[i] = W @ x_batch[i].

Inputs are float32 JAX arrays. A batch with no examples returns shape (0, d_out). Despite the function name, this exercise has no bias term: compute the linear map W @ x only.

Things to watch

in_axes=(0, 0) maps over rows of both arguments. In the worked example, their mapped sizes are three and two, so JAX raises an error. Even if those sizes match, each row of W would be paired with one example, producing dot products rather than the requested matrix-vector results.

None applies to the whole W argument here. It means no mapped axis for that argument, not that one of its dimensions should be removed. The mapping in this exercise uses positional arguments; keyword arguments passed to a vmapped function are mapped over their leading axis.

The controls below illustrate the worked example's axis placement. Changing them does not recompute the matrix product.

Loading visualization…