Skip to content

← JAX Vectorization & Control Flow step 25 of 25

Medium Primitives

lax.map vs vmap

Jump to the task

Apply a function to one slice at a time

lax.map(f, xs) applies f to slices along the leading axis and stacks the results. With its default batch_size=None, it processes elements sequentially using JAX control flow. vmap instead transforms the computation into batched operations. Sequential here describes the map across slices, not every operation inside f: a matrix operation in f can still use parallel hardware. This is not a Python loop that calls your function once per element at runtime.

Sequential mapping can reduce temporary working memory compared with a fully vectorized computation. The input and stacked output still take memory proportional to their sizes; total memory is not constant. Performance depends on the workload, so measure rather than assuming one is always faster.

lax.map also supports a batch_size option for processing chunks with vmap. This exercise uses the default element-by-element form.

Worked mini-example

from jax import lax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
out = lax.map(lambda xi: xi * xi, x)
# → [1.0, 4.0, 9.0]

The function f receives a single element (same shape as xs[0]), not the full batch.

For a matrix, a slice is a whole row. For example, lax.map(lambda row: row.sum(), jnp.array([[1., 2.], [3., 4.]])) returns [3., 7.]. The input has shape (2, 2), each row has shape (2,), each sum is scalar, and stacking the sums gives shape (2,). Map does not automatically visit every scalar in a multidimensional array.

Stage in the square example Shape
Input x (3,)
Each xi and its square scalar ()
Stacked result (3,)

Choose based on the workload

Squaring needs little temporary storage, so it does not demonstrate a memory advantage for lax.map. The distinction matters more when each slice creates large intermediate arrays. With a positive batch_size, map groups slices into vmapped chunks; f still describes one slice. A final incomplete chunk is handled separately.

Keep f pure: return its result instead of appending to a Python list or changing a global variable. Python side effects happen during tracing and are not a reliable record of per-slice execution. Unlike scan, map does not pass a changing carry from one slice to the next.

Problem

Implement lax_map_square(x) that squares each element using lax.map. Pass a helper or lambda that returns xi * xi (or xi ** 2) for one scalar. For this exercise, use the default map rather than squaring the whole input directly or using vmap.

The input is a one-dimensional real floating-point JAX array. Return its element-wise square with the same shape and dtype. Negative values become nonnegative; an empty input returns an empty array.

Output tests check the result, but cannot enforce your choice of API or establish a memory or speed advantage.

See the official lax.map reference.