jax-ml/jax · error · ValueError

All arrays must have matching shapes along non-concatenated

Error message

All arrays must have matching shapes along non-concatenated dimensions, got shape {arr.shape} at index {i} (expected dim {d} to be {arr0.shape[d]})

What it means

Like numpy.concatenate, all dimensions except the concatenation axis must match exactly across arrays; there is no broadcasting in Mosaic's register-level concat. The loop checks every non-axis dim and raises ValueError on the first mismatch, reporting the array index and dim.

Source

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

  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"
            f" {d} to be {arr0.shape[d]})"
        )
    new_shape[axis] += arr.shape[axis]
  new_shape = tuple(new_shape)

  match arr0.layout:
    case TiledLayout():
      for i, arr in enumerate(arrays[1:], start=1):
        if arr.layout != arr0.layout:
          raise ValueError(
              f"All arrays must have the same layout, got {arr.layout} at"
              f" index {i} (expected {arr0.layout})"
          )
      new_regs = np.concatenate([arr.registers for arr in arrays], axis=axis)
      return FragmentedArray(
          _registers=new_regs, _layout=arr0.layout, _is_signed=arr0.is_signed

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Reshape/pad the fragments so all non-concat dims match arrays[0].shape.
  2. Double-check which axis you actually split along and pass that as the concat axis.
  3. Pre-validate shapes: expected = arrays[0].shape; assert all(a.shape[:axis]+a.shape[axis+1:] == expected[:axis]+expected[axis+1:] for a in arrays).

Example fix

# before
out = FragmentedArray.concatenate([a, b], axis=0)  # a.shape (4,8), b.shape (4,16)
# after
b_padded = pad_fragment(b, target_shape=a.shape)  # make dim 1 match
out = FragmentedArray.concatenate([a, b_padded], axis=0)
Defensive patterns

Strategy: validation

Validate before calling

a0 = arrays[0].shape
for i, a in enumerate(arrays[1:], 1):
    for d in range(len(a0)):
        if d != axis and a.shape[d] != a0[d]:
            raise ValueError(f'shape mismatch at index {i}, dim {d}')
out = FragmentedArray.concatenate(arrays, axis=axis)

Prevention

When it happens

Trigger: Concatenating fragments whose shapes agree on the concat axis but differ elsewhere, e.g. (4, 8) and (4, 16) along axis=0 (dim 1 differs).

Common situations: Assuming numpy-style broadcasting carries over to fragments; splits along the wrong axis; tile shapes differing across pipeline stages (e.g., different vec_size or padded tiles).

Related errors

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