jax-ml/jax · error · NotImplementedError

Slicing batch dimensions is not supported.

Error message

Slicing batch dimensions is not supported.

What it means

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').

Source

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

  def commute_ndindexer(
      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer
  ) -> tuple[indexing.NDIndexer, state_types.Transform]:
    del aval
    batch_shape = self.batch_shape
    k = len(batch_shape)
    if len(indexer.indices) != k + 2:
      raise ValueError(
          f"Expected indexer to have exactly {k + 2} dimensions, "
          f"but got {len(indexer.indices)}."
      )
    batch_indices = indexer.indices[:-2]
    row_idx = indexer.indices[-2]
    col_idx = indexer.indices[-1]

    for idx in batch_indices:
      if isinstance(idx, indexing.Slice):
        raise NotImplementedError("Slicing batch dimensions is not supported.")

    batch_size = math.prod(batch_shape)
    m, n = indexer.shape[-2], indexer.shape[-1]
    physical_shape = (m, batch_size * n)

    batch_idx = 0
    for idx, size in zip(batch_indices, batch_shape):
      assert isinstance(idx, indexing.IntIndexer)
      batch_idx = batch_idx * size + idx

    if isinstance(col_idx, indexing.Slice):
      # We shift the column slice by batch_idx * n.
      new_col_idx = indexing.Slice(
          batch_idx * n + col_idx.start, col_idx.size, col_idx.stride
      )
    else:
      new_col_idx = batch_idx * n + col_idx

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index batch dims with concrete integers, not slices
2. Loop over the batch indices you need, selecting one at a time
3. Use a non-batch-folded layout if you need batch slicing

Example fix

// before
sub = ref[0:2, :, :]  # slice on batch dim -> NotImplementedError
// after
sub0 = ref[0, :, :]
sub1 = ref[1, :, :]

for i in indices[:-2]:
    assert not isinstance(i, indexing.Slice), 'batch dims require concrete indices'

When it happens

Trigger: Using an indexing.Slice (rather than an int/indexer) on any of the k batch dimensions of a ref under ExpandLeadingBatchDimensionsTransform.

Common situations: 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.

Related errors

• Non-slice indices are not supported in 2 minormost dims: {id
• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Explicit sharding is not currently supported in the pallas-t
• group_offset is not currently supported in the pallas-triton
• get not supported yet

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