jax-ml/jax · error · NotImplementedError

Scope {spec.trace_scope} not supported

Error message

Scope {spec.trace_scope} not supported

What it means

When generating the profiling instrumentation, the active trace scope must be WARP or WARPGROUP to determine which thread writes the trace (warp_idx vs warpgroup_idx). Any other scope on the spec raises NotImplementedError at instrumentation build time.

Source

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

      self,
      spec: ProfilerSpec,
      smem_buffer: ir.Value,
      gmem_buffer: ir.Value,
      wrap_in_custom_primitive: bool,
  ):
    i32 = ir.IntegerType.get_signless(32)
    index = ir.IndexType.get()
    self.spec = spec
    self.entries_per_wg = spec.entries_per_warpgroup
    self.wrap_in_custom_primitive = wrap_in_custom_primitive
    if spec.trace_scope == ThreadSubset.WARP:
      trace_idx = warp_idx(sync=False)
      scope_size = WARP_SIZE
    elif spec.trace_scope == ThreadSubset.WARPGROUP:
      trace_idx = warpgroup_idx(sync=False)
      scope_size = WARPGROUP_SIZE
    else:
      raise NotImplementedError(f"Scope {spec.trace_scope} not supported")
    trace_offset = arith.index_cast(
        index, arith.muli(trace_idx, c(self.entries_per_wg, i32))
    )
    smem_buffer = memref_slice(smem_buffer, ds(trace_offset, self.entries_per_wg))
    is_profiling_thread = arith.cmpi(
        arith.CmpIPredicate.eq,
        arith.remui(thread_idx(), c(scope_size, i32)),
        c(0, i32),
    )
    # Hopefully mem2reg will remove the allocation.
    offset = memref.alloca(ir.MemRefType.get((), index), [], [])
    memref.store(c(0, index), offset, [])
    self.ctx = _ProfilerCtx(
        start=globaltimer("low"),
        is_profiling_thread=is_profiling_thread,
        smem_buffer=smem_buffer,
        gmem_buffer=gmem_buffer,
        offset=offset,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct a fresh ProfilerSpec with WARP or WARPGROUP scope
2. Do not mutate spec.trace_scope after construction

When it happens

Trigger: Building a profiled kernel with a spec whose trace_scope is not WARP/WARPGROUP (again usually via post-construction mutation, since __init__ validates).

Common situations: Sharing/mutating a ProfilerSpec across differently-scoped builds.

Related errors

• {iterations=} must be positive
• Unsupported trace scope: {trace_scope}
• 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/6d4c39cc4f01c0c0. Report an issue: GitHub.