jax-ml/jax · error · ValueError

All arrays must have the same dtype, got {arr.mlir_dtype} at

Error message

All arrays must have the same dtype, got {arr.mlir_dtype} at index {i} (expected {arr0.mlir_dtype})

What it means

concatenate performs np.concatenate on the underlying register arrays, which requires a single MLIR element type. If any fragment's mlir_dtype differs from arrays[0], it raises ValueError naming the index and types.

Source

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

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Convert all fragments to a common dtype before concatenating (re-load/store or cast through registers/GMEM with the target element type).
Audit where each fragment is created and pin the dtype explicitly instead of relying on inference.
Pre-validate: assert all(a.mlir_dtype == arrays[0].mlir_dtype for a in arrays).

Example fix

# before
out = FragmentedArray.concatenate([a_f32, b_f16], axis=0)
# after
b_f32 = cast_fragment(b_f16, a_f32.mlir_dtype)
out = FragmentedArray.concatenate([a_f32, b_f32], axis=0)

Defensive patterns

Strategy: validation

Validate before calling

dt = arrays[0].mlir_dtype
assert all(a.mlir_dtype == dt for a in arrays), [
    (i, a.mlir_dtype) for i, a in enumerate(arrays)
]
out = FragmentedArray.concatenate(arrays, axis=axis)

Prevention

Pin dtypes explicitly when creating each fragment.
Convert to a common dtype before concat.
Watch f16/bf32 accumulators mixed with stored precisions.

When it happens

Trigger: Concatenating an f32 FragmentedArray with an f16 or i32 one, or mixing bf16/f32 fragments produced by different pipeline stages.

Common situations: Precision conversions (e.g., accumulating in f32 but storing fragments as f16/bf16); mixing dtypes across pipeline stages or after a cast op; dtype defaults changing between JAX versions.

Related errors

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