jax-ml/jax · error · ValueError

Unsupported shape: {x.shape}

Error message

Unsupported shape: {x.shape}

What it means

The inverse transform (undoing batch-dim expansion) requires the physical aval to be exactly 2-D, since the layout maps a (batch..., m, n) logical shape to a flat (m, batch*n) physical shape. ndim != 2 raises ValueError('Unsupported shape').

Source

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

@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the physical aval 2-D (m, batch*n) before the transform
2. Verify batch_shape passed to the transform matches how the buffer was allocated
3. Reconstruct the buffer with the correct 2-D physical layout

Example fix

// before
phys = ShapedArray((b, m, n), dt)  # ndim 3 -> ValueError
// after
phys = ShapedArray((m, b * n), dt)

assert x.ndim == 2, f'physical aval must be 2-D, got shape {x.shape}'

When it happens

Trigger: Applying the batch-expansion undo/transform to a ShapedArray with ndim other than 2 — e.g. trying to interpret a 3-D or 1-D physical buffer as an expanded-batch layout.

Common situations: Mismatch between the shape of a physical buffer and the declared batch_shape in a Mosaic kernel; hand-constructed avals passed to get_ref_aval/to_block_mapping.

Related errors

• Unsupported ndim: {x.ndim}
• only 1-dimensional input supported.
• 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/3d82e2b0948a9d58. Report an issue: GitHub.