{"record":{"id":"90c987035616c61d","repo":"jax-ml/jax","slug":"slicing-batch-dimensions-is-not-supported","errorCode":null,"errorMessage":"Slicing batch dimensions is not supported.","messagePattern":"Slicing batch dimensions is not supported\\.","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1350,"sourceCode":"\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)\n      batch_idx = batch_idx * size + idx\n\n    if isinstance(col_idx, indexing.Slice):\n      # We shift the column slice by batch_idx * n.\n      new_col_idx = indexing.Slice(\n          batch_idx * n + col_idx.start, col_idx.size, col_idx.stride\n      )\n    else:\n      new_col_idx = batch_idx * n + col_idx\n","sourceCodeStart":1332,"sourceCodeEnd":1368,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1332-L1368","documentation":"Batch-expansion commuting maps each batch index to an offset in the flattened column dimension, which requires batch indices to be concrete integers. indexing.Slice in any batch dimension raises NotImplementedError('Slicing batch dimensions is not supported').","triggerScenarios":"Using an indexing.Slice (rather than an int/indexer) on any of the k batch dimensions of a ref under ExpandLeadingBatchDimensionsTransform.","commonSituations":"Writing kernel code that slices a batch of matrices (e.g. ref[0:2, :, :]) when the batch dims have been folded into the column extent; treating folded-batch buffers like normal batched arrays.","solutions":["Index batch dims with concrete integers, not slices","Loop over the batch indices you need, selecting one at a time","Use a non-batch-folded layout if you need batch slicing"],"exampleFix":"// before\nsub = ref[0:2, :, :]  # slice on batch dim -> NotImplementedError\n// after\nsub0 = ref[0, :, :]\nsub1 = ref[1, :, :]","handlingStrategy":"validation","validationCode":"for i in indices[:-2]:\n    assert not isinstance(i, indexing.Slice), 'batch dims require concrete indices'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use integers on batch dims of batch-folded refs","Loop over batch indices instead of slicing"],"tags":["jax","pallas","indexing","slicing","not-implemented"],"backgroundTag":"unsupported-batch-slicing","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}