jax-ml/jax · error · ValueError

Second dimension {x.shape[1]} must be divisible by batch_siz

Error message

Second dimension {x.shape[1]} must be divisible by batch_size {batch_size}

What it means

When expanding batch dimensions, the second (column) dimension of the 2-D physical array must be divisible by the product of batch_shape, because each logical row holds batch*n elements. A column extent that isn't a multiple raises this ValueError.

Source

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

class ExpandLeadingBatchDimensionsTransform(state_types.Transform):
  """The inverse of CollapseLeadingBatchDimensionsTransform.

  Specifically, it maps `(m, math.prod(batch_shape) * n)` to `(*batch_shape, m,
  n)`.
  """

  batch_shape: tuple[int, ...] = jax.tree.static()

  def transform_type(
      self, x: jax_core.AbstractValue
  ) -> state_types.AbstractRef:
    match x:
      case jax_core.ShapedArray():
        if x.ndim != 2:
          raise ValueError(f"Unsupported shape: {x.shape}")
        batch_size = math.prod(self.batch_shape)
        if x.shape[1] % batch_size != 0:
          raise ValueError(
              f"Second dimension {x.shape[1]} must be divisible by batch_size"
              f" {batch_size}"
          )
        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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Size the physical second dim as batch_size * n
2. Pad n so shape[1] is a multiple of batch_size
3. Correct the batch_shape so it divides the column extent

Example fix

// before
phys = ShapedArray((m, 64), dt)  # batch_shape=(3,) -> 64 % 3 != 0
// after
phys = ShapedArray((m, 66), dt)  # 66 = 3 * 22

import math
bs = math.prod(batch_shape)
assert shape[1] % bs == 0, f'{shape[1]} not divisible by batch_size {bs}'

When it happens

Trigger: Passing a 2-D aval whose shape[1] % prod(batch_shape) != 0 to ExpandLeadingBatchDimensionsTransform.transform_type (e.g. batch_shape (3,) with a shape[1] of 64).

Common situations: Allocating physical buffers sized without accounting for the batch multiplier; irregular batch sizes (non-power-of-2 pipelines) in WGMMA layouts.

Related errors

• Block size must be a multiple of the input size. Got block {
• Every block dimension must be either a multiple or factor of
• Seed key_data must be 1D.
• Leading dimension of seed key_data must be 1.
• Not implemented: bitcast 1D

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