jax-ml/jax · error · NotImplementedError

Scope {self.trace_scope} not supported

Error message

Scope {self.trace_scope} not supported

What it means

A defensive NotImplementedError inside ProfilerSpec._num_traces: after __init__ validates the scope, any scope other than WARP/WARPGROUP reaching here means an inconsistent state (e.g. a manually mutated spec.trace_scope).

Source

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

    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, ...]
  ) -> ir.MemRefType:
    return ir.MemRefType.get(
        (self._num_traces(grid, block) * self.entries_per_warpgroup,),
        ir.IntegerType.get_signless(32),
    )

  def jax_buffer_type(
      self, grid: tuple[int, ...], block: tuple[int, ...]
  ) -> jax.ShapeDtypeStruct:
    return jax.ShapeDtypeStruct(
        (self._num_traces(grid, block) * self.entries_per_warpgroup,),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create a new ProfilerSpec with the desired scope instead of mutating trace_scope
2. Never mutate spec.trace_scope after construction

When it happens

Trigger: Assigning spec.trace_scope = ThreadSubset.BLOCK after construction and then calling spec.dump(...), spec.mlir_buffer_type(...), or spec.smem_i32_elements(...).

Common situations: Mutable spec objects shared between kernels; bypassing the constructor check by attribute mutation.

Related errors

• Scope {self.spec.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/3bc9bfc2cc24b684. Report an issue: GitHub.