Skip to content

← JAX Profiling & Debugging step 3 of 12

Easy Primitives

@profiler.annotate_function

Jump to the task

Label a function call in a profile

jax.profiler.annotate_function wraps a Python function call in a TraceAnnotation event. During an active profiling capture, the event marks the interval in which that wrapper executes. The decorator does not start a capture or change the function's mathematical result.

This differs from jax.named_scope, which adds names to staged JAX operations. A host event and compiled-operation metadata answer different questions about where time is spent.

import jax
import jax.numpy as jnp
from functools import partial

@partial(jax.profiler.annotate_function, name="sum_squares")
def compute(x):
    return jnp.sum(x ** 2)

result = compute(jnp.array([3., 4.]))  # scalar 25

The name override is optional. Without it, JAX uses the function's qualified name when available, so a nested function or method can include its enclosing function or class. Closures and methods are supported; they do not require switching to named_scope.

What does the interval include?

A JAX array may return before its device computation finishes. This wrapper records the host call, not necessarily the entire device execution. Waiting inside the wrapper would extend the event through completion, but also changes synchronization and potentially performance.

Decorator order matters with JIT. A profiler wrapper outside a jitted helper runs on each Python call. An annotation inside a jitted body runs when Python traces that body, not on every cached execution. Keep the host wrapper outside JIT when you want per-call host events.

Problem

Implement annotated_call(x) as an ordinary Python wrapper:

  1. Define an inner function named inner and decorate it with @jax.profiler.annotate_function.
  2. Have inner return jnp.sum(x ** 2).
  3. Call inner(x) and return its result.

X is a one-dimensional float32 JAX array. Return a scalar JAX array. For [3, 4], return 9 + 16 = 25; an empty array returns zero. Do not start a profiling capture or add a timing return value in the submission. The exercise labels a call without requiring a wait.

Output tests verify the sum, not the presence or name of a trace event. Inspect the decorator and use a separate profile to examine events.