{"record":{"id":"a5b3d82ad4949624","repo":"jax-ml/jax","slug":"all-arrays-must-have-matching-shapes-along-non-con","errorCode":null,"errorMessage":"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]})","messagePattern":"All arrays must have matching shapes along non-concatenated dimensions, got shape (.+?) at index (.+?) \\(expected dim (.+?) to be (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5370,"sourceCode":"  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      )\n    if arr.is_signed != arr0.is_signed:\n      raise ValueError(\n          f\"All arrays must have the same signedness, got {arr.is_signed} at\"\n          f\" index {i} (expected {arr0.is_signed})\"\n      )\n    for d in range(rank):\n      if d != axis and arr.shape[d] != arr0.shape[d]:\n        raise ValueError(\n            \"All arrays must have matching shapes along non-concatenated\"\n            f\" dimensions, got shape {arr.shape} at index {i} (expected dim\"\n            f\" {d} to be {arr0.shape[d]})\"\n        )\n    new_shape[axis] += arr.shape[axis]\n  new_shape = tuple(new_shape)\n\n  match arr0.layout:\n    case TiledLayout():\n      for i, arr in enumerate(arrays[1:], start=1):\n        if arr.layout != arr0.layout:\n          raise ValueError(\n              f\"All arrays must have the same layout, got {arr.layout} at\"\n              f\" index {i} (expected {arr0.layout})\"\n          )\n      new_regs = np.concatenate([arr.registers for arr in arrays], axis=axis)\n      return FragmentedArray(\n          _registers=new_regs, _layout=arr0.layout, _is_signed=arr0.is_signed","sourceCodeStart":5352,"sourceCodeEnd":5388,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5352-L5388","documentation":"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.","triggerScenarios":"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).","commonSituations":"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).","solutions":["Reshape/pad the fragments so all non-concat dims match arrays[0].shape.","Double-check which axis you actually split along and pass that as the concat axis.","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)."],"exampleFix":"# before\nout = FragmentedArray.concatenate([a, b], axis=0)  # a.shape (4,8), b.shape (4,16)\n# after\nb_padded = pad_fragment(b, target_shape=a.shape)  # make dim 1 match\nout = FragmentedArray.concatenate([a, b_padded], axis=0)","handlingStrategy":"validation","validationCode":"a0 = arrays[0].shape\nfor i, a in enumerate(arrays[1:], 1):\n    for d in range(len(a0)):\n        if d != axis and a.shape[d] != a0[d]:\n            raise ValueError(f'shape mismatch at index {i}, dim {d}')\nout = FragmentedArray.concatenate(arrays, axis=axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Remember Mosaic concat never broadcasts.","Split and re-join along the same axis.","Pad/reshape fragments so non-concat dims match."],"tags":["jax","mosaic-gpu","validation","shape-mismatch"],"backgroundTag":"shape-mismatch-on-concat","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}