jax-ml/jax · error · ValueError

Expected a single dimension when passing a single index

Error message

Expected a single dimension when passing a single index

What it means

Raised in cluster_idx (utils.py:2362) when the caller passes a single scalar dim_idx (an ir.Value) while specifying multiple dims. A scalar index can only broadcast against one dimension, so the combination is ambiguous.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2362

def cluster_idx(
    dim: gpu.Dimension | Sequence[gpu.Dimension] | None = None,
    dim_idx: ir.Value | Sequence[ir.Value] | None = None,
) -> ir.Value:
  """Returns the linear index of a block within a subset of the cluster spanned by the given dimensions.

  dim_idx can be used to specify the index of another block along the selected
  dimensions. If not provided, the current block's index is used.
  """
  if dim is None:
    dim = tuple(gpu.Dimension)
  elif isinstance(dim, gpu.Dimension):
    dim = (dim,)
  if dim_idx is None:
    dim_idx = [gpu.cluster_block_id(d) for d in dim]
  elif isinstance(dim_idx, ir.Value):
    if len(dim) != 1:
      raise ValueError(
          "Expected a single dimension when passing a single index"
      )
    dim_idx = [dim_idx]
  index = ir.IndexType.get()
  stride = c(1, index)
  lin_idx = c(0, index)
  for d, idx in sorted(zip(dim, dim_idx, strict=True), key=lambda x: x[0]):
    lin_idx = arith.addi(lin_idx, arith.muli(idx, stride))
    stride = arith.muli(stride, gpu.cluster_dim_blocks(d))
  return lin_idx


def get_cluster_ptr(
    ptr: ir.Value, cluster_block: ir.Value, generic: bool = True
):
  i32 = ir.IntegerType.get_signless(32)
  assert cluster_block.type == i32, cluster_block.type
  assert ptr.type == llvm.PointerType.get(3), ptr.type

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass dim_idx as a sequence matching len(dim), e.g. [idx_x, idx_y]
2. Or pass dim as a single gpu.Dimension if only one index is available

Example fix

# before
cluster_idx(dim=(gpu.Dimension.x, gpu.Dimension.y), dim_idx=idx)
# after
cluster_idx(dim=(gpu.Dimension.x, gpu.Dimension.y), dim_idx=[idx_x, idx_y])

idx = dim_idx if isinstance(dim_idx, (list, tuple)) else ([dim_idx] if dim_idx is not None else None)
assert idx is None or len(idx) == len(dim)

try:
    cluster_idx(dim=dim, dim_idx=dim_idx)
except ValueError as e:
    if 'single dimension' in str(e):
        dim_idx = [dim_idx] * len(dim)
    else:
        raise

When it happens

Trigger: Calling utils.cluster_idx(dim=(gpu.Dimension.x, gpu.Dimension.y), dim_idx=some_value) — more than one dim but dim_idx is a single ir.Value rather than a sequence.

Common situations: Refactoring code that indexed one dimension and adding another dim without turning dim_idx into a list; copy-pasting multi-dim calls from single-dim examples.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• dims and idxs must have the same length
• Cluster barriers are not arrays
• cluster= must be at most 3D, got {self}.

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