jax-ml/jax · error · ValueError

Unsupported trace scope: {trace_scope}

Error message

Unsupported trace scope: {trace_scope}

What it means

ProfilerSpec.__init__ validates that trace_scope is one of ThreadSubset.WARP or ThreadSubset.WARPGROUP, since trace buffer layout depends on the scope size (32 vs 128 threads). Any other value (including other ThreadSubset enum members like BLOCK) is rejected.

Source

Thrown at jax/experimental/mosaic/gpu/profiler.py:203

        event costs 2 entries, and 3 entries are reserved for a header.
      dump_path: Where to write the trace.
      trace_scope: Whether one trace covers a warp or a warpgroup.
      bounds_check: If True, events past the buffer capacity are dropped (the
        trace is truncated) at the cost of a slightly higher per-event overhead.
        If False (default), overflowing the buffer corrupts neighbouring SMEM,
        which usually crashes the kernel.
    """
    self.entries_per_warpgroup = entries_per_warpgroup
    self.interned_names: dict[str, int] = {}
    self.bounds_check = bounds_check
    if dump_path == "sponge":
      self.dump_path = os.getenv(
          "TEST_UNDECLARED_OUTPUTS_DIR", tempfile.gettempdir()
      )
    else:
      self.dump_path = dump_path
    if trace_scope not in (ThreadSubset.WARP, ThreadSubset.WARPGROUP):
      raise ValueError(f"Unsupported trace scope: {trace_scope}")
    self.trace_scope = trace_scope

  def _num_traces(
      self, grid: tuple[int, ...], block: tuple[int, ...]
  ) -> int:
    if self.trace_scope == ThreadSubset.WARP:
      scope_size = WARP_SIZE
    elif self.trace_scope == ThreadSubset.WARPGROUP:
      scope_size = WARPGROUP_SIZE
    else:
      raise NotImplementedError(f"Scope {self.trace_scope} not supported")

    if math.prod(block) % scope_size:
      raise ValueError(f"Block size is not a multiple of {scope_size}")
    return math.prod(grid) * math.prod(block) // scope_size

  def mlir_buffer_type(
      self, grid: tuple[int, ...], block: tuple[int, ...]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use ThreadSubset.WARP or ThreadSubset.WARPGROUP explicitly
2. If you need block-level tracing, profile per warp and aggregate in dump()

Example fix

# before
spec = ProfilerSpec(trace_scope=ThreadSubset.BLOCK)

# after
spec = ProfilerSpec(trace_scope=ThreadSubset.WARPGROUP)

from jax.experimental.mosaic.gpu import profiler as P
assert P.ProfilerSpec.trace_scope.__class__  # inspect
valid = {P.ThreadSubset.WARP, P.ThreadSubset.WARPGROUP}
if trace_scope not in valid:
    raise ValueError(f'trace_scope must be one of {valid}')

def is_supported_scope(s) -> bool:
    return s in (ThreadSubset.WARP, ThreadSubset.WARPGROUP)

When it happens

Trigger: Constructing ProfilerSpec(trace_scope=ThreadSubset.BLOCK) or passing an arbitrary integer/string as trace_scope.

Common situations: Assuming the profiler works at block granularity because ThreadSubset has other members; upgrading code written against a version that silently accepted other scopes.

Related errors

• {iterations=} must be positive
• Scope {spec.trace_scope} not supported
• Expected both or neither of scales to be specified.
• Scope {self.trace_scope} not supported
• Block size is not a multiple of {scope_size}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/2b22b1c8bc96f048. Report an issue: GitHub.