jax-ml/jax · error · ValueError

{axis=} is out of bounds for array of {rank=}

Error message

{axis=} is out of bounds for array of {rank=}

What it means

FragmentedArray.concatenate bounds-checks the axis argument like numpy: it must satisfy -rank <= axis < rank. Since FragmentedArray has no dynamic axes, an out-of-range axis cannot be resolved and is rejected with ValueError.

Source

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

      pre_indices.append(i)
  # If both exist, all expansions must happen before all preserved
  # dimensions.
  if exp_indices and pre_indices and max(exp_indices) >= min(pre_indices):
    return False
  return True


def concatenate(
    arrays: Sequence[FragmentedArray],
    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})"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a negative axis or clamp: axis = axis % rank (Python modulo handles negatives).
2. Check the rank first: assert -rank <= axis < rank, or derive axis from len(arr.shape).
3. Verify intermediate ops (squeeze/reshape) produced the rank you expect before concatenating.

Example fix

# before
out = FragmentedArray.concatenate(arrs, axis=2)  # rank-2 arrays -> error
# after
out = FragmentedArray.concatenate(arrs, axis=1)

rank = len(arrays[0].shape)
if not -rank <= axis < rank:
    raise ValueError(f'bad axis {axis} for rank {rank}')
axis = axis % rank
out = FragmentedArray.concatenate(arrays, axis=axis)

When it happens

Trigger: Passing axis >= rank or axis < -rank, e.g. axis=2 for a 2D fragmented array, or axis=-3 on a rank-2 array.

Common situations: Copy-pasting numpy code with a hardcoded axis onto lower-rank fragments; computing axis from another array's rank; off-by-one mistakes after squeezing/expanding dims.

Related errors

• 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}.
• cluster_names must have the same length as cluster, got {sel

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