jax-ml/jax · error · ValueError

dims and idxs must have the same length

Error message

dims and idxs must have the same length

What it means

ClusterRefTransform validates that dims (cluster axis names) and idxs (block indices) tuples have equal length; mismatched lengths raise ValueError in the dataclass's __post_init__.

Source

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

    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):
      raise TypeError("ref must be a reference")
    ref = pallas_core.TransformedRef(ref, transforms=())
  if any(isinstance(t, MulticastRef) for t in ref.transforms):
    raise ValueError("Can't make a multicast reference into a peer reference.")
  return pallas_core.TransformedRef(
      ref.ref, (*ref.transforms, PeerMemRef(device_id, device_id_type)),
  )


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ClusterRefTransform(state_types.Transform):
  dims: tuple[jax_core.AxisName, ...] = jax.tree.static()
  idxs: tuple[Any, ...]

  def __post_init__(self):
    if len(self.dims) != len(self.idxs):
      raise ValueError("dims and idxs must have the same length")

  def transform_type(self, x):
    return x

  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    raise NotImplementedError()

  def commute_ndindexer(
      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, ClusterRefTransform]:
    return indexer, self


def cluster_ref(
    ref: _Ref,
    block_id: dict[jax_core.AxisName, Any],
) -> pallas_core.TransformedRef:
  """Translate memref to a peer memref in the cluster."""

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the dims and idxs tuples are built from the same dict (zip of keys and values)
2. Validate lengths before constructing: len(dims) == len(idxs)
3. Prefer using cluster_ref(ref, block_id) which builds the tuples consistently

Example fix

// before
ClusterRefTransform(('x', 'y'), (0,))

// after
ClusterRefTransform(('x', 'y'), (0, 1))

assert len(dims) == len(idxs), 'dims/idxs length mismatch'

When it happens

Trigger: Constructing ClusterRefTransform(dims, idxs) directly with tuples of different lengths, or calling cluster_ref with a dict whose keys/values get split inconsistently.

Common situations: Manually building cluster reference transforms; bugs in code generating block_id mappings for cluster kernels.

Related errors

• cluster= must be at most 3D, got {self}.
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• grid_names must have the same length as grid, got {self}.
• cluster_names must have the same length as cluster, got {sel
• num_threads and thread_name must be either both set or both

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