Skip to content

← JAX Foundations step 7 of 25

Medium Primitives

2-D Scatter via .at[].add()

Jump to the task

Add values to grid cells

Suppose several measurements belong to the same cell in a grid. To find each cell's total, start with zeros and add each measurement at its row and column. This is a 2-D scatter-add, useful for weighted histograms.

out.at[rows, cols].add(values) pairs the indices by position: values[k] goes to cell (rows[k], cols[k]). It does not update every combination of the supplied rows and columns.

Worked example

import jax.numpy as jnp

grid = jnp.zeros((3, 3))
rows = jnp.array([0, 1, 0])
cols = jnp.array([1, 1, 1])
values = jnp.array([1.0, 2.0, 3.0])
result = grid.at[rows, cols].add(values)
# [[0., 4., 0.],
#  [0., 2., 0.],
#  [0., 0., 0.]]

The first and third values both go to (0, 1), so that cell holds 1 + 3 = 4. The second value goes to (1, 1). All other cells stay zero. The update returns a new array value; grid still contains zeros.

Problem

Implement scatter_2d_indices(out_shape, row_indices, col_indices, values). Create a zero-filled array of shape (H, W) and add each value at its paired row and column. Return the resulting array. Repeated pairs add together, and empty update arrays leave the output full of zeros.

Inputs are float32 JAX arrays:

  • out_shape has shape (2,) and contains whole-number dimensions [H, W], both positive.
  • row_indices, col_indices, and values have shape (N,).
  • Indices are whole numbers within the output bounds and fit in int32. Convert both index arrays to integer dtype before indexing.
  • Values are finite and may be positive, zero, or negative.

For example, shape [2, 3], rows [0, 1], columns [2, 0], and values [7, 3] give [[0, 0, 7], [3, 0, 0]].

Shape handling and JIT

For this exercise's ordinary, eager call, use int(out_shape[0]) and int(out_shape[1]) to obtain the dimensions for jnp.zeros((H, W)). This conversion does not make a dynamic shape work under jax.jit: converting a traced array value to a Python integer raises an error.

If you adapt this function for JIT, supply the shape as a Python tuple of integers and mark that argument static, or use a fixed output shape inside a compiled helper. Marking the original JAX array static is not a substitute: JAX arrays are not hashable static arguments. The scatter operation itself supports JIT when the output shape is known.

Things to watch

  • Casting indices changes their dtype; it does not validate them. Fractional or out-of-range indices are outside this exercise's contract.
  • Use .add, not .set: replacing a cell would lose contributions.
  • Repeated updates all contribute, but their floating-point addition order is not guaranteed. Last-bit rounding can vary across backends.
  • Do not rely on an indexing exception to catch invalid data. By default, out-of-bounds scatter updates are dropped, and negative indices count from the end. Explicit mode="clip" instead clamps indices to the bounds.