Skip to content

← JAX Vectorization & Control Flow step 1 of 25

Medium Primitives

vmap with in_axes

Jump to the task

Why this matters

jax.vmap turns a function for one example into a function for a batch. in_axes tells it which axis to map in each positional argument. The function sees a slice with that axis removed.

For a matrix of rows, in_axes=0 supplies one row at a time conceptually. JAX transforms the array operations rather than running a Python loop for every row.

Worked mini-example

import jax
import jax.numpy as jnp

a_batch = jnp.array([[1.0, 2.0], [3.0, 4.0]])
b_batch = jnp.array([[5.0, 6.0], [7.0, 8.0]])
pair_outer = jax.vmap(jnp.outer, in_axes=(0, 0))
result = pair_outer(a_batch, b_batch)
# result[0] = [[5, 6], [10, 12]]
# result[1] = [[21, 24], [28, 32]]
# result.shape = (2, 2, 2)

Each corresponding pair of length-two rows produces a two-by-two matrix. The default output axis stacks those matrices along axis 0. This computes two pairs, not all four combinations of rows.

Common pitfalls

  • All mapped axes must have the same length. A mapped length-one axis is not automatically broadcast to a longer mapped axis.
  • None means an argument is shared without being sliced. It does not make the argument static for JIT.
  • For two array arguments, in_axes=(0, 0) and in_axes=0 are equivalent. If you provide a tuple, it needs one entry per positional argument.
  • Pass inputs positionally when using this tuple. Keyword arguments to a vmapped function are mapped over axis 0.
  • Mapping axis 1 of a matrix supplies columns, not rows.

Problem

Implement batched_outer(a_batch, b_batch) with jax.vmap and explicit in_axes=(0, 0). Do not use a loop.

Both inputs have shape (N, d). Return shape (N, d, d), with out[i, j, k] = a_batch[i, j] * b_batch[i, k]. This is an outer product for each row pair, not a dot product.

The diagram follows the worked example: two row pairs produce two matrices, stacked into shape (2, 2, 2).

Loading visualization…