{"record":{"id":"a8ba3998f8d0e870","repo":"jax-ml/jax","slug":"axis-is-out-of-bounds-for-array-of-rank","errorCode":null,"errorMessage":"{axis=} is out of bounds for array of {rank=}","messagePattern":"(.+?) is out of bounds for array of (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5344,"sourceCode":"      pre_indices.append(i)\n  # If both exist, all expansions must happen before all preserved\n  # dimensions.\n  if exp_indices and pre_indices and max(exp_indices) >= min(pre_indices):\n    return False\n  return True\n\n\ndef concatenate(\n    arrays: Sequence[FragmentedArray],\n    axis: int = 0,\n) -> FragmentedArray:\n  \"\"\"Concatenates fragmented arrays along the specified axis.\"\"\"\n  if not arrays:\n    raise ValueError(\"Need at least one array to concatenate\")\n  arr0 = arrays[0]\n  rank = len(arr0.shape)\n  if not -rank <= axis < rank:\n    raise ValueError(f\"{axis=} is out of bounds for array of {rank=}\")\n  if axis < 0:\n    axis += rank\n\n  if len(arrays) == 1:\n    return arr0\n\n  new_shape = list(arr0.shape)\n  for i, arr in enumerate(arrays[1:], start=1):\n    if len(arr.shape) != rank:\n      raise ValueError(\n          f\"All arrays must have the same rank, got {len(arr.shape)} at index\"\n          f\" {i} (expected {rank})\"\n      )\n    if arr.mlir_dtype != arr0.mlir_dtype:\n      raise ValueError(\n          f\"All arrays must have the same dtype, got {arr.mlir_dtype} at\"\n          f\" index {i} (expected {arr0.mlir_dtype})\"\n      )","sourceCodeStart":5326,"sourceCodeEnd":5362,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5326-L5362","documentation":"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.","triggerScenarios":"Passing axis >= rank or axis < -rank, e.g. axis=2 for a 2D fragmented array, or axis=-3 on a rank-2 array.","commonSituations":"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.","solutions":["Use a negative axis or clamp: axis = axis % rank (Python modulo handles negatives).","Check the rank first: assert -rank <= axis < rank, or derive axis from len(arr.shape).","Verify intermediate ops (squeeze/reshape) produced the rank you expect before concatenating."],"exampleFix":"# before\nout = FragmentedArray.concatenate(arrs, axis=2)  # rank-2 arrays -> error\n# after\nout = FragmentedArray.concatenate(arrs, axis=1)","handlingStrategy":"validation","validationCode":"rank = len(arrays[0].shape)\nif not -rank <= axis < rank:\n    raise ValueError(f'bad axis {axis} for rank {rank}')\naxis = axis % rank\nout = FragmentedArray.concatenate(arrays, axis=axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Normalize axis with axis % rank before passing.","Log rank alongside axis in kernel authoring helpers.","Re-check rank after squeeze/reshape steps."],"tags":["jax","mosaic-gpu","validation","axis-bounds"],"backgroundTag":"axis-out-of-bounds","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}