Skip to content

← JAX Profiling & Debugging step 2 of 12

Medium Primitives

profiler.StepTraceAnnotation

Jump to the task

Mark iterations in a profile

jax.profiler.StepTraceAnnotation creates a profiler event for a logical step, such as one iteration of a training loop. The name identifies the activity; step_num distinguishes iterations. It does not run an optimizer, change the result or start a capture.

import jax
import jax.numpy as jnp
from jax import profiler as jpr

@jax.jit
def step_sum(x):
    return jnp.sum(x)

x = jnp.array([1., 2., 3.])
step_sum(x).block_until_ready()  # warm up separately
with jpr.StepTraceAnnotation("train", step_num=0):
    result = step_sum(x)
    result.block_until_ready()
# result is 6

During an active capture, the profiler can use the event to organize per-step analysis. The wait in this example keeps the host interval open until the result is ready. Without it, the context can end while device work remains pending. Adding a wait at every step changes execution overlap, so choose it deliberately for your measurement.

Keep step metadata on the host

Put the annotation around calls to a compiled step in the Python loop. An annotation inside a jitted body executes during tracing, so it does not create a fresh host event for each cached call.

Here step arrives as a whole Python number supplied as a float. Convert it to int for the metadata. int() does not make a traced value concrete: applying JIT to this whole wrapper with a dynamic step can fail while tracing. Keep this host wrapper outside JIT instead of treating the step number as numerical input to the compiled calculation.

Problem

Implement annotated_step(x, step):

  1. Enter jpr.StepTraceAnnotation("train", step_num=int(step)).
  2. Inside the context, compute and return jnp.sum(x).

X is a one-dimensional float32 JAX array. Step is a nonnegative whole Python number supplied as a float. Return a scalar JAX array: [1, 2, 3] sums to 6, [-1, 1] sums to zero, and an empty input also sums to zero. Changing step changes metadata, not the sum.

The task does not require a synchronization wait or a profiling capture. It demonstrates annotation syntax; it does not promise completed device work at context exit. Output tests check numerical values, not whether a step event was captured. Inspect the context and its metadata too.