Skip to content

← JAX Vectorization & Control Flow step 2 of 25

Medium Primitives

vmap with out_axes

Jump to the task

Why this matters

in_axes selects the input axis that contains examples. out_axes chooses where the mapped axis appears in the result. These choices are independent.

If one example returns a vector of length d, mapping over N examples gives shape (N, d) with out_axes=0, or (d, N) with out_axes=1. If one example returns a scalar, the mapped result is a vector with just one axis.

Worked mini-example

import jax
import jax.numpy as jnp

matrix = jnp.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
column_sums = jax.vmap(
    lambda col: jnp.sum(col ** 2),
    in_axes=1,
    out_axes=0,
)
result = column_sums(matrix)  # [17, 29, 45], shape (3,)

The first column is [1, 4], so its result is 1 + 16 = 17. Each column produces one number; stacking three numbers gives a length-three vector.

To see output-axis placement with vector results, compare:

rows = jax.vmap(lambda col: col * 2, in_axes=1, out_axes=0)(matrix)
columns = jax.vmap(lambda col: col * 2, in_axes=1, out_axes=1)(matrix)
# rows.shape = (3, 2); columns.shape = (2, 3)
# rows.T and columns contain the same values.

Common pitfalls

  • Selecting input axis 1 does not select output axis 1. The default out_axes remains zero.
  • For the scalar result in this exercise, out_axes=0 or -1 is valid. Axis 1 is out of range for the result vector.
  • in_axes=0 would sum squared rows instead, giving a different result. Use a rectangular example to make that mistake visible.
  • These axis choices describe logical array dimensions, not a guarantee about physical memory layout.

Problem

Implement vmap_with_out_axis(matrix) to compute the sum of squares of each column. Use jax.vmap with explicit in_axes=1 and out_axes=0.

The input has shape (d, N). Return shape (N,), where out[j] = sum(matrix[:, j] ** 2).

With no columns, the output is empty. With zero rows, each column has an empty sum and contributes zero.

In the diagram, mapping axis 1 of [2, 3] supplies length-two columns. Each result is a scalar ([]); stacking gives [3].

Loading visualization…