Skip to content

← JAX Profiling & Debugging step 1 of 12

Medium Primitives

jax.named_scope for Profiler Grouping

Jump to the task

Attach meaningful names to JAX operations

A compiled calculation can contain many operations. Names such as forward_pass or normalization help connect those operations to the Python calculation that produced them.

jax.named_scope adds a name to JAX's name stack. Operations created inside the scope carry that context into staged code and profiling metadata. It does not change their mathematical result or start a profiling capture.

import jax
import jax.numpy as jnp

@jax.jit
def sum_sq(x):
    with jax.named_scope("forward_pass"):
        return jnp.sum(x ** 2)

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

The square and reduction are associated with forward_pass when JAX traces them. Scopes can nest, so operation names can include several levels of context. The supplied name is part of that context, not a promise that every viewer will show a separate timeline span with exactly that text. Compiler transformations and viewer presentation affect what you see.

Operation names versus host events

named_scope can be used inside JIT because it supplies metadata during tracing. TraceAnnotation and annotate_function instead create profiler events around Python execution. Use those when the host interval is what you want to mark. None of these tools automatically waits for asynchronous device work or measures a duration for your return value.

This exercise uses the context-manager form, with a with block. named_scope can also decorate a function; it is not restricted to context-manager syntax.

Problem

Implement named_scope_compute(x):

  1. Enter with jax.named_scope("forward_pass"):.
  2. Compute and return jnp.sum(x ** 2) inside that block.

X is a one-dimensional float32 JAX array. Return a scalar JAX array. For [3, 4], the result is 25. For an empty input, it is zero. Return the calculation, not the context manager or the scope name.

Numerical tests verify the sum. They cannot prove that the named scope is present, so inspect the block as well. A profile or staged-program inspection is needed to examine the operation metadata.