jax-ml/jax · error · NotImplementedError

Scope {self.spec.trace_scope} not supported

Error message

Scope {self.spec.trace_scope} not supported

What it means

During finalize() (called from _launch), the code computes the global trace index from block/warp or warpgroup indices; a scope other than WARP/WARPGROUP has no defined trace layout and raises NotImplementedError before emitting the final copy of traces to global memory.

Source

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

    index = ir.IndexType.get()
    i32 = ir.IntegerType.get_signless(32)

    with self._profiler_ctx() as ctx:
      gpu.barrier()  # Make sure all warpgroups are done.

      block_idx = c(0, index)
      for dim in gpu.Dimension:
        block_idx = arith.addi(
            arith.muli(block_idx, gpu.grid_dim(dim)), gpu.block_id(dim)
        )
      if self.spec.trace_scope == ThreadSubset.WARP:
        trace_idx = warp_idx(sync=False)
        traces_per_block = math.prod(block) // WARP_SIZE
      elif self.spec.trace_scope == ThreadSubset.WARPGROUP:
        trace_idx = warpgroup_idx(sync=False)
        traces_per_block = math.prod(block) // WARPGROUP_SIZE
      else:
        raise NotImplementedError(f"Scope {self.spec.trace_scope} not supported")
      global_trace_idx = arith.addi(
          arith.muli(block_idx, c(traces_per_block, index)),
          arith.index_cast(index, trace_idx),
      )
      start_offset = arith.muli(global_trace_idx, c(self.entries_per_wg, index))
      wg_gmem_buffer = memref_slice(
          ctx.gmem_buffer, ds(start_offset, self.entries_per_wg)
      )
      with when(ctx.is_profiling_thread):
        memref.store(ctx.start, wg_gmem_buffer, [c(0, index)])
        memref.store(smid(), wg_gmem_buffer, [c(1, index)])
        num_traces = arith.index_cast(i32, memref.load(ctx.offset, []))
        memref.store(num_traces, wg_gmem_buffer, [c(2, index)])
        traces = vector.load(
            ir.VectorType.get((self.entries_per_wg - 3,), i32),
            ctx.smem_buffer,
            [c(0, index)],
        )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Recreate the ProfilerSpec with a supported scope
2. Treat ProfilerSpec as immutable

When it happens

Trigger: A spec with an unsupported trace_scope reaching kernel launch/finalize — practically only through attribute mutation after construction.

Common situations: Long-lived mutable spec objects reused across launches.

Related errors

• Scope {self.trace_scope} not supported
• {iterations=} must be positive
• Unsupported trace scope: {trace_scope}
• Block size is not a multiple of {scope_size}
• Allocated too many names

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