{"record":{"id":"1b6f0af803d04c16","repo":"jax-ml/jax","slug":"expected-indexer-to-have-exactly-k-2-dimension","errorCode":null,"errorMessage":"Expected indexer to have exactly {k + 2} dimensions, but got {len(indexer.indices)}.","messagePattern":"Expected indexer to have exactly (.+?) dimensions, but got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1340,"sourceCode":"          )\n        transformed_shape = self.batch_shape + (\n            x.shape[0],\n            x.shape[1] // batch_size,\n        )\n        return x.update(shape=transformed_shape)\n      case state_types.AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case _:\n        raise TypeError(f\"Unsupported type: {x}\")\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, state_types.Transform]:\n    del aval\n    batch_shape = self.batch_shape\n    k = len(batch_shape)\n    if len(indexer.indices) != k + 2:\n      raise ValueError(\n          f\"Expected indexer to have exactly {k + 2} dimensions, \"\n          f\"but got {len(indexer.indices)}.\"\n      )\n    batch_indices = indexer.indices[:-2]\n    row_idx = indexer.indices[-2]\n    col_idx = indexer.indices[-1]\n\n    for idx in batch_indices:\n      if isinstance(idx, indexing.Slice):\n        raise NotImplementedError(\"Slicing batch dimensions is not supported.\")\n\n    batch_size = math.prod(batch_shape)\n    m, n = indexer.shape[-2], indexer.shape[-1]\n    physical_shape = (m, batch_size * n)\n\n    batch_idx = 0\n    for idx, size in zip(batch_indices, batch_shape):\n      assert isinstance(idx, indexing.IntIndexer)","sourceCodeStart":1322,"sourceCodeEnd":1358,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1322-L1358","documentation":"commute_ndindexer for the batch-expansion transform requires the indexer's rank to be exactly k+2, where k = len(batch_shape): k batch indices plus one row and one column index. Any other number of indices raises ValueError with both expected and actual counts.","triggerScenarios":"Indexing an expanded-batch ref with fewer or more indices than batch_shape rank + 2 — e.g. dropping the batch dims (ref[m_idx, n_idx] when batch_shape has rank 1).","commonSituations":"Indexing a logically (batch..., m, n) buffer as if it were 2-D; mismatch between the grid/batch shape declared for the kernel and the indexer used in kernel body code.","solutions":["Supply exactly k+2 indices, including full batch-dim indices (e.g. slices covering each batch dim)","Check len(batch_shape) of the transform and match your indexer rank to it","Fix the batch_shape declaration if it doesn't match the logical rank you index with"],"exampleFix":"// before\nsub = ref[i_m, i_n]  # 2 indices, batch_shape rank 1 -> ValueError\n// after\nsub = ref[:, i_m, i_n]  # 3 = 1 + 2 indices","handlingStrategy":"validation","validationCode":"k = len(batch_shape)\nassert len(indices) == k + 2, f'need {k+2} indices, got {len(indices)}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Index with exactly batch_rank + 2 indices","Keep grid/batch_shape declarations in sync with kernel indexing code"],"tags":["jax","pallas","indexing","rank-mismatch"],"backgroundTag":"indexer-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}