Dynamic Slice

Move the window, keep its size fixed

Suppose you need three neighboring values from a vector, but the start position is only known when the function runs. lax.dynamic_slice lets that position change under JIT while keeping the output shape fixed. For a one-dimensional array, pass one start index and one window size:

import jax
import jax.numpy as jnp
from jax import lax

@jax.jit
def three_values(x, start):
    return lax.dynamic_slice(x, (start,), (3,))

x = jnp.arange(10, dtype=jnp.float32)
three_values(x, 2)  # [2.0, 3.0, 4.0]
three_values(x, 5)  # [5.0, 6.0, 7.0]

Each start is an integer scalar, one per input axis. It does not have to be int32, though this exercise uses that conversion. An arbitrary array of indices is not a single start position.

Static size means known during tracing

Ordinary slicing such as x[start:start + 3] fails under JIT when start is traced. Dynamic slicing allows the start to vary, but its window size must still be known during tracing. Calling int(size) does not make a traced size static; it raises ConcretizationTypeError.

For the function you will write, mark window_size as static explicitly:

compiled = jax.jit(extract_window, static_argnames=("window_size",))
compiled(x, 2.0, window_size=3.0)  # [2.0, 3.0, 4.0]
compiled(x, 5.0, window_size=3.0)  # [5.0, 6.0, 7.0]

Pass a Python number for that static argument, not a JAX array. Changing the size can require a new compilation because the result's shape changes. The float representation here is part of the exercise interface; sizes and starts must have whole-number values.

At the boundary, the window shifts

The result keeps its requested length. For x = [0, 1, ..., 9] and size 3, start 9 returns [7, 8, 9], not a shortened [9] and not wrapped values [9, 0, 1]. The start is adjusted so the window fits.

Negative starts are interpreted relative to the end by default, then adjusted if necessary. Start -1 also returns [7, 8, 9]; start -4 returns [6, 7, 8]. A start below -len(x) is adjusted to the beginning. This is not circular slicing or validation of the requested position.

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.