{"record":{"id":"d0cd1325a29fa335","repo":"jax-ml/jax","slug":"all-arrays-must-have-the-same-dtype-got-arr-mlir","errorCode":null,"errorMessage":"All arrays must have the same dtype, got {arr.mlir_dtype} at index {i} (expected {arr0.mlir_dtype})","messagePattern":"All arrays must have the same dtype, got (.+?) at index (.+?) \\(expected (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5359,"sourceCode":"  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      )\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","sourceCodeStart":5341,"sourceCodeEnd":5377,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5341-L5377","documentation":"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.","triggerScenarios":"Concatenating an f32 FragmentedArray with an f16 or i32 one, or mixing bf16/f32 fragments produced by different pipeline stages.","commonSituations":"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.","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)."],"exampleFix":"# before\nout = FragmentedArray.concatenate([a_f32, b_f16], axis=0)\n# after\nb_f32 = cast_fragment(b_f16, a_f32.mlir_dtype)\nout = FragmentedArray.concatenate([a_f32, b_f32], axis=0)","handlingStrategy":"validation","validationCode":"dt = arrays[0].mlir_dtype\nassert all(a.mlir_dtype == dt for a in arrays), [\n    (i, a.mlir_dtype) for i, a in enumerate(arrays)\n]\nout = FragmentedArray.concatenate(arrays, axis=axis)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pin dtypes explicitly when creating each fragment.","Convert to a common dtype before concat.","Watch f16/bf32 accumulators mixed with stored precisions."],"tags":["jax","mosaic-gpu","validation","dtype-mismatch"],"backgroundTag":"dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}