{"record":{"id":"d00e1be3318b238a","repo":"jax-ml/jax","slug":"all-wgstridedfraglayout-arrays-must-have-the-same","errorCode":null,"errorMessage":"All WGStridedFragLayout arrays must have the same vec_size, got {arr.layout.vec_size} at index {i} (expected {vec_size})","messagePattern":"All WGStridedFragLayout arrays must have the same vec_size, got (.+?) at index (.+?) \\(expected (.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5403,"sourceCode":"          )\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}\"\n          )\n        if arr.layout.vec_size != vec_size:\n          raise ValueError(\n              \"All WGStridedFragLayout arrays must have the same vec_size,\"\n              f\" got {arr.layout.vec_size} at index {i} (expected {vec_size})\"\n          )\n      new_layout = WGStridedFragLayout(shape=new_shape, vec_size=vec_size)\n      new_regs = np.concatenate([arr.registers for arr in arrays], axis=0)\n      return FragmentedArray(\n          _registers=new_regs, _layout=new_layout, _is_signed=arr0.is_signed\n      )\n\n    case WGSplatFragLayout():\n      raise NotImplementedError(\n          \"Concatenating arrays with splat layout is not supported.\"\n      )\n\n    case layout:\n      assert_never(layout)\n","sourceCodeStart":5385,"sourceCodeEnd":5420,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5385-L5420","documentation":"Mosaic GPU's FragmentedArray.concat requires all concatenated arrays that use WGStridedFragLayout to share the same vec_size (vector width). The loop validates each array's layout against the first array's vec_size and raises ValueError on mismatch. This invariant is required because a single WGStridedFragLayout with one vec_size is constructed for the concatenated result.","triggerScenarios":"Calling FragmentedArray.concat (or APIs building on it) where the first array has a WGStridedFragLayout with vec_size=N but a later array (index i) has vec_size != N, e.g. mixing arrays produced with different element vectorization (vec2 vs vec4 register fragments).","commonSituations":"Combining tensors materialized with different wgmma/mma vector widths, upgrading/downgrading Mosaic versions that changed default vec_size, or manually constructing FragmentedArray instances with mismatched layouts before concatenation.","solutions":["Make all inputs use the same vec_size: convert/reshape fragments (e.g. via layout conversion ops) so every WGStridedFragLayout array matches the first array's vec_size before concat","Check arr.layout.vec_size for every input before calling concat and re-emit the ones that differ","If mixing widths is intentional, concatenate via materializing to registers/tensors and re-fragmenting instead of FragmentedArray.concat"],"exampleFix":"# before\nresult = FragmentedArray.concat([a, b])  # a.vec_size=2, b.vec_size=4 -> ValueError\n\n# after\nassert all(arr.layout.vec_size == arrays[0].layout.vec_size for arr in arrays)\nresult = FragmentedArray.concat([a, b])  # both vec_size=4","handlingStrategy":"validation","validationCode":"from jax.experimental.mosaic.gpu import fragmented_array as fa\n\ndef can_concat(arrays):\n  first = arrays[0].layout\n  if isinstance(first, fa.WGStridedFragLayout):\n    return all(\n        isinstance(a.layout, fa.WGStridedFragLayout)\n        and a.layout.vec_size == first.vec_size\n        for a in arrays\n    )\n  return True","typeGuard":"def same_vec_size(arrays) -> bool:\n  vs = getattr(arrays[0].layout, 'vec_size', None)\n  return vs is None or all(getattr(a.layout, 'vec_size', None) == vs for a in arrays)","tryCatchPattern":"try:\n  out = FragmentedArray.concat(arrays)\nexcept ValueError as e:\n  if 'vec_size' in str(e):\n    raise RuntimeError(f'Incompatible vec_size: {[a.layout for a in arrays]}') from e\n  raise","preventionTips":["Assert vec_size equality across fragment arrays before concat","Avoid mixing fragments produced under different vectorization settings in one kernel","Log arr.layout for each input when building kernels that concatenate fragments"],"tags":["mosaic-gpu","fragmented-array","layout-mismatch","concat","validation"],"backgroundTag":"layout-validation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}