take_along_axis (Gather)

Why this matters

Sometimes each row of a matrix needs a different set of columns. jnp.take_along_axis(matrix, indices, axis=1) uses each row of indices to choose values from the matching matrix row.

For example, after sorting scores separately in each row, you can use the resulting indices to gather the scores in that order.

Worked mini-example

import jax.numpy as jnp

matrix = jnp.array([[10, 20, 30, 40], [50, 60, 70, 80]])
indices = jnp.array([[3, 0], [2, 1]], dtype=jnp.int32)
out = jnp.take_along_axis(matrix, indices, axis=1)
# out = [[40, 10], [70, 60]], shape (2, 2)

Row zero selects columns 3 and 0. Row one selects columns 2 and 1. The order of indices becomes the order of values in the result. Repeating an index repeats the selected value.

Common pitfalls

• With an explicit axis, the input and indices need the same number of dimensions. Other dimensions must be compatible for broadcasting. This exercise uses the same row count.
• Indices must have an integer dtype. The fixtures supply integer-valued floats, so convert them with indices.astype(jnp.int32).
• jnp.take(matrix, [3, 0], axis=1) selects columns 3 and 0 in every row. It flattens the input only when axis is omitted. It does not apply a separate index row to each matrix row.
• Do not rely on an exception for an invalid index. JAX gather behavior depends on the operation and its bounds mode. The inputs for this exercise stay within the valid range.
• Casting is not validation: a fractional input would be truncated. This task supplies whole-number indices.

Select values using coordinates

A slice selects a regularly spaced region of an array. A gather selects entries using supplied indices. Check whether those indices select whole rows or columns, or a different entry within each row.

For [[10, 11], [20, 21]], choosing column 1 from the first row and column 0 from the second gives [11, 20]. Choosing columns [1, 0] from every row instead gives [[11, 10], [21, 20]]. These operations use the same index values but have different meanings and output shapes.

Scatter writes values to selected destinations

Scatter takes values and destination indices. The operation may replace existing values or combine them using a reduction such as addition.

Starting with [0, 0, 0], scatter-add of values [2, 3] at indices [1, 1] gives [0, 5, 0]. Both values contribute to the repeated destination. Repeated assignment is not equivalent to this sum; follow the operation's rule for repeated indices.

Selecting with a mask differs from replacing with a mask

For values [10, 20, 30] and mask [True, False, True], selecting masked entries gives [10, 30]. Replacing masked entries with zero gives [0, 20, 0]. Selection can change the shape; replacement preserves it. A selection whose size depends on runtime values needs particular care in compiled code with fixed output shapes.

Check whether an update changes the original array

JAX indexed updates return a new array value. In PyTorch, basic slicing can return a view that shares storage, while advanced indexing returns a copy. Modifying a selected copy does not modify the source. Direct indexed assignment is a separate operation, so distinguish changing a selection from assigning into the original.