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

Validation in ClusterRefTransform.__post_init__ (Mosaic GPU): the cluster dims and per-dimension indices must correspond one-to-one so each axis has an index; unequal lengths raise. The inputs at fault are the dims and idxs fields of the cluster ref transform.

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
• Can't instantiate with arguments.
• cluster_names must have the same length as cluster, got
• Expected metadata dtype to be uint2, got
• grid_names must have the same length as grid, got

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

Appendix: 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)