{"record":{"id":"1cfe4dba0262c2be","repo":"jax-ml/jax","slug":"all-arrays-must-have-the-same-layout-got-arr-lay","errorCode":null,"errorMessage":"All arrays must have the same layout, got {arr.layout} at index {i} (expected {arr0.layout})","messagePattern":"All arrays must have the same layout, got (.+?) at index (.+?) \\(expected (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5382,"sourceCode":"      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\n      )\n\n    case WGStridedFragLayout(vec_size=vec_size):\n      if axis != 0:\n        raise NotImplementedError(\n            \"Concatenating arrays with strided layout is only supported along\"\n            \" axis 0\"\n        )\n      for i, arr in enumerate(arrays[1:], start=1):\n        if not isinstance(arr.layout, WGStridedFragLayout):\n          raise ValueError(\n              f\"Expected WGStridedFragLayout, got {arr.layout} at index {i}\"","sourceCodeStart":5364,"sourceCodeEnd":5400,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5364-L5400","documentation":"In the TiledLayout branch, concatenate verifies every array has the identical layout because np.concatenate on registers is only layout-preserving when tiles, register arrangement, and swizzle all match. A differing TiledLayout raises ValueError with both layouts printed.","triggerScenarios":"Concatenating fragments built with different tiled layouts, e.g. different tile shapes, different register orderings, or fragments produced by load_tiled with different swizzle values.","commonSituations":"Producing fragments in different passes or with different tiling parameters (tile shape, swizzle) and then combining them; JAX version changes altering default tiled layouts; mixing fragments from GMEM loads with fragments constructed directly from registers.","solutions":["Re-materialize the mismatched fragments through the same layout as arrays[0] (store to GMEM/SMEM and reload with arrays[0].layout parameters).","Ensure all fragments are created by the same load/config path so their TiledLayout compares equal.","Print and compare .layout of each array before concat to find the divergent one early."],"exampleFix":"# before\nout = FragmentedArray.concatenate([a, b], axis=0)  # b.layout != a.layout\n# after\nb_relaid = b.to_tiled_layout_like(a)  # store + reload with a's tiling params\nout = FragmentedArray.concatenate([a, b_relaid], axis=0)","handlingStrategy":"validation","validationCode":"lay = arrays[0].layout\nbad = [i for i, a in enumerate(arrays) if a.layout != lay]\nassert not bad, f'layout mismatch at {bad}: {lay}'\nout = FragmentedArray.concatenate(arrays, axis=axis)","typeGuard":"from jax.experimental.mosaic.gpu.fragmented_array import TiledLayout\n\ndef all_same_tiled_layout(arrays) -> bool:\n    return isinstance(arrays[0].layout, TiledLayout) and all(\n        a.layout == arrays[0].layout for a in arrays\n    )","tryCatchPattern":null,"preventionTips":["Produce all concatenated fragments via the same load path and tiling params.","Compare .layout before concatenating in debug prints.","Round-trip through memory to normalize layouts."],"tags":["jax","mosaic-gpu","validation","layout-mismatch"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}