jax-ml/jax · error · ValueError

num_threads and thread_name must be either both set or both

Error message

num_threads and thread_name must be either both set or both None, got {self}

What it means

Launch-config validation: `num_threads` and `thread_name` are an optional pair and must be set together or both left as None. Setting only one would create a named warp/thread axis with no size (or vice versa), so the dataclass rejects it.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1650

  cluster_names: Sequence[str] = ()
  # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.
  num_threads: int | None = None
  thread_name: str | None = None
  kernel_name: str | None = None

  def __post_init__(self):
    if len(self.cluster) > 3:
      raise ValueError(f"cluster= must be at most 3D, got {self}.")
    if len(self.grid_names) != len(self.grid):
      raise ValueError(
          f"grid_names must have the same length as grid, got {self}."
      )
    if len(self.cluster_names) != len(self.cluster):
      raise ValueError(
          f"cluster_names must have the same length as cluster, got {self}."
      )
    if (self.thread_name is None) != (self.num_threads is None):
      raise ValueError(
          "num_threads and thread_name must be either both set or both None,"
          f" got {self}"
      )
    max_mosaic_threads = 2048 // 128
    if self.num_threads is not None and self.num_threads > max_mosaic_threads:
      raise ValueError(
          "Requested too many CUDA threads per block. Each Mosaic thread"
          f" corresponds to 128 CUDA threads. At most {max_mosaic_threads}"
          f" are supported, got {self}"
      )
    object.__setattr__(self, "grid", tuple(self.grid))
    object.__setattr__(self, "grid_names", tuple(self.grid_names))
    object.__setattr__(self, "cluster", tuple(self.cluster))
    object.__setattr__(self, "cluster_names", tuple(self.cluster_names))

  @property
  def default_memory_space(self) -> MemorySpace:
    return MemorySpace.GMEM

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Set both: num_threads=4, thread_name='warp'
Or remove both to disable explicit threading

Example fix

// before
num_threads=4
thread_name=None

// after
num_threads=4
thread_name='warp'

Defensive patterns

Strategy: validation

Validate before calling

assert (num_threads is None) == (thread_name is None), 'set both or neither'

Type guard

def threading_pair_valid(num_threads, thread_name) -> bool:
    return (num_threads is None) == (thread_name is None)

Try / catch

try:
    cfg = LaunchConfig(num_threads=n, thread_name=t)
except ValueError:
    if n is not None and t is None:
        t = 'warp'
    cfg = LaunchConfig(num_threads=n, thread_name=t)

Prevention

Treat num_threads/thread_name as one setting; set them together in a helper
Disable both when not using warp specialization

When it happens

Trigger: Constructing the launch config with num_threads=4 but thread_name=None, or thread_name='warp' without num_threads.

Common situations: Enabling warp specialization by adding num_threads but forgetting the name used by warp-specialization checks inside the kernel; refactoring away thread names while leaving num_threads behind.

Related errors

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