{"record":{"id":"f33255d71958a382","repo":"jax-ml/jax","slug":"dims-and-idxs-must-have-the-same-length","errorCode":null,"errorMessage":"dims and idxs must have the same length","messagePattern":"dims and idxs must have the same length","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1070,"sourceCode":"    if not isinstance(jax_core.typeof(ref), state_types.AbstractRef):\n      raise TypeError(\"ref must be a reference\")\n    ref = pallas_core.TransformedRef(ref, transforms=())\n  if any(isinstance(t, MulticastRef) for t in ref.transforms):\n    raise ValueError(\"Can't make a multicast reference into a peer reference.\")\n  return pallas_core.TransformedRef(\n      ref.ref, (*ref.transforms, PeerMemRef(device_id, device_id_type)),\n  )\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass ClusterRefTransform(state_types.Transform):\n  dims: tuple[jax_core.AxisName, ...] = jax.tree.static()\n  idxs: tuple[Any, ...]\n\n  def __post_init__(self):\n    if len(self.dims) != len(self.idxs):\n      raise ValueError(\"dims and idxs must have the same length\")\n\n  def transform_type(self, x):\n    return x\n\n  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:\n    raise NotImplementedError()\n\n  def commute_ndindexer(\n      self, _: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, ClusterRefTransform]:\n    return indexer, self\n\n\ndef cluster_ref(\n    ref: _Ref,\n    block_id: dict[jax_core.AxisName, Any],\n) -> pallas_core.TransformedRef:\n  \"\"\"Translate memref to a peer memref in the cluster.\"\"\"","sourceCodeStart":1052,"sourceCodeEnd":1088,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1052-L1088","documentation":"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__.","triggerScenarios":"Constructing ClusterRefTransform(dims, idxs) directly with tuples of different lengths, or calling cluster_ref with a dict whose keys/values get split inconsistently.","commonSituations":"Manually building cluster reference transforms; bugs in code generating block_id mappings for cluster kernels.","solutions":["Ensure the dims and idxs tuples are built from the same dict (zip of keys and values)","Validate lengths before constructing: len(dims) == len(idxs)","Prefer using cluster_ref(ref, block_id) which builds the tuples consistently"],"exampleFix":"// before\nClusterRefTransform(('x', 'y'), (0,))\n\n// after\nClusterRefTransform(('x', 'y'), (0, 1))","handlingStrategy":"validation","validationCode":"assert len(dims) == len(idxs), 'dims/idxs length mismatch'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Build dims/idxs from the same dict via keys()/values()"],"tags":["jax","pallas","mosaic-gpu","cluster","validation"],"backgroundTag":"argument-length-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}