jax-ml/jax · error · ValueError

Unsupported ndim: {x.ndim}

Error message

Unsupported ndim: {x.ndim}

What it means

ExpandLeadingBatchDimensionsTransform.transform_type requires the input ShapedArray to have ndim >= 2, because it folds all leading dims into a batch factor of a 2-D (rows x cols) physical layout. Arrays with 0 or 1 dims cannot be represented and raise ValueError('Unsupported ndim').

Source

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

    return pp.text(f"{{unswizzle({self.swizzle})}}")


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class CollapseLeadingBatchDimensionsTransform(state_types.Transform):
  """A transform that collapses leading batch dimensions into the minor dimension.

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

  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 ndim: {x.ndim}")
        batch_size = math.prod(x.shape[:-2])
        transformed_shape = (x.shape[-2], batch_size * x.shape[-1])
        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 undo(self, x: jax_core.AbstractValue) -> state_types.Transform:
    assert hasattr(x, "shape")
    return ExpandLeadingBatchDimensionsTransform(x.shape[:-2])


@tree_util.register_dataclass
@dataclasses.dataclass(frozen=True)
class ExpandLeadingBatchDimensionsTransform(state_types.Transform):
  """The inverse of CollapseLeadingBatchDimensionsTransform.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Reshape to at least 2-D before applying the transform (e.g. x[None, :])
2. Store 1-D data in a separate non-transformed buffer
3. Pad the vector to a (1, n) matrix

Example fix

// before
ref = make_ref(vec)  # vec.ndim == 1 -> ValueError
// after
ref = make_ref(vec[None, :])  # shape (1, n)

assert x.ndim >= 2, f'need ndim >= 2, got {x.ndim}'

def is_batch_expandable(x) -> bool:\n    return getattr(x, 'ndim', 0) >= 2

When it happens

Trigger: Passing a scalar (ndim 0) or 1-D array through a transform that expands leading batch dimensions, e.g. mapping a 1-D vector ref to a 2-D swizzled buffer via get_ref_aval/to_block_mapping.

Common situations: Trying to store 1-D bias/scale vectors in a 2-D-only WGMMA layout inside Mosaic kernels; allocating scalars in transforms meant for matrices.

Related errors

• Unsupported shape: {x.shape}
• 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/a18a297235e06f5b. Report an issue: GitHub.