Skip to content

← JAX Vectorization & Control Flow step 19 of 25

Medium Primitives

jit of scan

Jump to the task

Compile the work around a scan

lax.scan already compiles its loop body. Wrapping a larger function in jit can also compile work around that loop, such as preprocessing inputs and reducing its outputs. A Python loop calling a jitted step is another valid design when the host needs to handle loading or logging between steps.

Repeated compatible calls can reuse a compiled executable. Different shapes, dtypes or static argument values can trigger recompilation; actual performance depends on the workload and hardware.

Worked mini-example

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

@jax.jit
def scan_product(x):
    def step(c, xi):
        new = c * xi
        return new, new
    _, ys = lax.scan(
        step, 1.0, x
    )
    return ys

x = jnp.array([1., 2., 3., 4.])
result = scan_product(x)
# [1., 2., 6., 24.]

Here the carry is a scalar product and the output contains every intermediate product. The exercise below asks for a sum and only its final carry instead.

@jax.jit decorates the function; compilation is lazy. Define and reuse the jitted function rather than recreating it inside each iteration of a host loop if you want its compilation cache to be useful.

Explore the running product

Step through the example's arithmetic below. This browser simulation shows carry and output values; it does not compile or benchmark JAX. The other operation buttons let you compare how the same scan structure behaves with a different body.

Loading visualization…

Common pitfalls

  • jit-ing the call, not the function. jax.jit(my_fn(x)) evaluates my_fn(x) eagerly and then tries to jit the result, not the function. Always write jax.jit(my_fn)(x) or decorate my_fn with @jax.jit.
  • scan body not pure. Any Python side effect inside step (print, list append) runs during tracing, not once per runtime step. Retracing can run it again. Keep scan bodies pure.
  • Confusing carry and outputs. lax.scan returns (final_carry, stacked_ys). If you only need the final accumulated value, take final_carry (index 0) and emit None as the per-step output.

Problem

Implement jit_scan_sum(x) by decorating the function with @jax.jit and using lax.scan inside it. The input is a one-dimensional real floating-point JAX array. Return a scalar with the same dtype containing the sum of all elements, not a vector of prefix sums. An empty input returns zero.

Start with jnp.zeros((), dtype=x.dtype) as the carry. Each step adds its input and emits None. Return the final carry from scan. Keep the decorated function at module scope so repeated compatible calls can reuse it. Output tests verify sums, not JIT or cache use.