jax-ml/jax · error · ValueError

Expected indexer to have exactly {k + 2} dimensions, but got

Error message

Expected indexer to have exactly {k + 2} dimensions, but got {len(indexer.indices)}.

What it means

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.

Source

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

          )
        transformed_shape = self.batch_shape + (
            x.shape[0],
            x.shape[1] // batch_size,
        )
        return x.update(shape=transformed_shape)
      case state_types.AbstractRef():
        return x.update(inner_aval=self.transform_type(x.inner_aval))
      case _:
        raise TypeError(f"Unsupported type: {x}")

  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Supply exactly k+2 indices, including full batch-dim indices (e.g. slices covering each batch dim)
2. Check len(batch_shape) of the transform and match your indexer rank to it
3. Fix the batch_shape declaration if it doesn't match the logical rank you index with

Example fix

// before
sub = ref[i_m, i_n]  # 2 indices, batch_shape rank 1 -> ValueError
// after
sub = ref[:, i_m, i_n]  # 3 = 1 + 2 indices

k = len(batch_shape)
assert len(indices) == k + 2, f'need {k+2} indices, got {len(indices)}'

When it happens

Trigger: 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).

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

Related errors

• Unsupported range type: {type(r)}.
• Cannot do int indexing on TPU
• Can only load a single key per load.
• Dimension must be 0 for 1D iota.
• Integer indexing of refs that follows a non-trivial slice is

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