{"record":{"id":"31d0970e534c7440","repo":"jax-ml/jax","slug":"expected-wgstridedfraglayout-got-arr-layout-at","errorCode":null,"errorMessage":"Expected WGStridedFragLayout, got {arr.layout} at index {i}","messagePattern":"Expected WGStridedFragLayout, got (.+?) at index (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5399,"sourceCode":"        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}\"\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","sourceCodeStart":5381,"sourceCodeEnd":5417,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5381-L5417","documentation":"In the WGStridedFragLayout branch, concatenate requires every array to also have a WGStridedFragLayout, because register concatenation relies on the shared strided structure. A TiledLayout (or other layout) mixed into the list raises ValueError with its index.","triggerScenarios":"Concatenating a WGMMA accumulator fragment (strided layout) with a fragment loaded from tiled SMEM/GMEM (TiledLayout) in the same list.","commonSituations":"Mixing WGMMA outputs with directly-loaded fragments and concatenating them; pipeline refactors that change how one fragment is materialized; fragments created via different load helpers with different layout inference.","solutions":["Convert all fragments to WGStridedFragLayout (or all to TiledLayout) before concatenating, e.g., round-trip through memory with a consistent load path.","Group fragments by layout type and concatenate within each group before combining.","Check isinstance(arr.layout, WGStridedFragLayout) per element before the call."],"exampleFix":"# before\nout = FragmentedArray.concatenate([wgmma_acc, tiled_frag], axis=0)  # error at index 1\n# after\ntiled_as_strided = reload_with_strided_layout(tiled_frag)\nout = FragmentedArray.concatenate([wgmma_acc, tiled_as_strided], axis=0)","handlingStrategy":"type-guard","validationCode":"from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout\n\nassert all(isinstance(a.layout, WGStridedFragLayout) for a in arrays), [\n    (i, type(a.layout).__name__) for i, a in enumerate(arrays)\n]\nout = FragmentedArray.concatenate(arrays, axis=0)","typeGuard":"from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout\n\ndef all_strided(arrays) -> bool:\n    return all(isinstance(a.layout, WGStridedFragLayout) for a in arrays)","tryCatchPattern":"try:\n    out = FragmentedArray.concatenate(arrays, axis=0)\nexcept ValueError as e:\n    if 'WGStridedFragLayout' not in str(e): raise\n    arrays = [to_strided(a) for a in arrays]  # normalize via memory round-trip\n    out = FragmentedArray.concatenate(arrays, axis=0)","preventionTips":["Don't mix WGMMA accumulators with tiled-loaded fragments in one concat.","Group fragments by layout type before combining.","Normalize layouts via store/reload when combining heterogeneous sources."],"tags":["jax","mosaic-gpu","validation","layout-mismatch","wgmma"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}