jax-ml/jax · error · ValueError

All arrays must have the same rank, got {len(arr.shape)} at

Error message

All arrays must have the same rank, got {len(arr.shape)} at index {i} (expected {rank})

What it means

concatenate requires all FragmentedArrays to share the same rank (number of dimensions) because registers are concatenated with np.concatenate, which needs uniformly shaped arrays apart from the concat axis. A mismatched rank raises ValueError with the offending index.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:5354

    axis: int = 0,
) -> FragmentedArray:
  """Concatenates fragmented arrays along the specified axis."""
  if not arrays:
    raise ValueError("Need at least one array to concatenate")
  arr0 = arrays[0]
  rank = len(arr0.shape)
  if not -rank <= axis < rank:
    raise ValueError(f"{axis=} is out of bounds for array of {rank=}")
  if axis < 0:
    axis += rank

  if len(arrays) == 1:
    return arr0

  new_shape = list(arr0.shape)
  for i, arr in enumerate(arrays[1:], start=1):
    if len(arr.shape) != rank:
      raise ValueError(
          f"All arrays must have the same rank, got {len(arr.shape)} at index"
          f" {i} (expected {rank})"
      )
    if arr.mlir_dtype != arr0.mlir_dtype:
      raise ValueError(
          f"All arrays must have the same dtype, got {arr.mlir_dtype} at"
          f" index {i} (expected {arr0.mlir_dtype})"
      )
    if arr.is_signed != arr0.is_signed:
      raise ValueError(
          f"All arrays must have the same signedness, got {arr.is_signed} at"
          f" index {i} (expected {arr0.is_signed})"
      )
    for d in range(rank):
      if d != axis and arr.shape[d] != arr0.shape[d]:
        raise ValueError(
            "All arrays must have matching shapes along non-concatenated"
            f" dimensions, got shape {arr.shape} at index {i} (expected dim"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Normalize ranks before concatenating: expand or squeeze dims so every array has rank equal to arrays[0].
2. Fix the producing op that introduced the extra/missing dimension (check layout/shape construction).
3. Add a pre-call assert: assert all(len(a.shape) == len(arrays[0].shape) for a in arrays).

Example fix

# before
out = FragmentedArray.concatenate([a, b], axis=0)  # a rank 2, b rank 3
# after
b2 = b.reshape_with_layout(...) # or store/load so ranks match
out = FragmentedArray.concatenate([a, b2], axis=0)

rank = len(arrays[0].shape)
bad = [i for i, a in enumerate(arrays) if len(a.shape) != rank]
assert not bad, f'rank mismatch at {bad}'
out = FragmentedArray.concatenate(arrays, axis=axis)

def all_same_rank(arrays) -> bool:
    r = len(arrays[0].shape)
    return all(len(a.shape) == r for a in arrays)

When it happens

Trigger: Mixing a 2D FragmentedArray with a 3D one in the arrays list, e.g. concatenate([a_2d, b_3d]).

Common situations: Some pipeline stages adding or removing a unit dimension (expand_dims/squeeze) so fragments drift in rank; heterogeneously-built fragments appended to one list; refactoring from numpy where broadcasting silently handled rank differences.

Related errors

• Length of major_to_minor and the rank of the value should ma
• dims and idxs must have the same length
• Swizzle {self.swizzle} is not supported. Only 32, 64 and 128
• cluster= must be at most 3D, got {self}.
• grid_names must have the same length as grid, got {self}.

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