{"record":{"id":"f7ff547c32d754a2","repo":"jax-ml/jax","slug":"concatenating-arrays-with-strided-layout-is-only-s","errorCode":null,"errorMessage":"Concatenating arrays with strided layout is only supported along axis 0","messagePattern":"Concatenating arrays with strided layout is only supported along axis 0","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5393,"sourceCode":"    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}\"\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      )","sourceCodeStart":5375,"sourceCodeEnd":5411,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5375-L5411","documentation":"For WGStridedFragLayout fragments, concatenate is implemented by appending whole register arrays (row-major register blocks per warpgroup), which only preserves the strided layout semantics when concatenating along axis 0. Any other axis raises NotImplementedError.","triggerScenarios":"Calling FragmentedArray.concatenate(arrs, axis=1) (or any nonzero axis) where arrays[0].layout is a WGStridedFragLayout (typical for WGMMA accumulator fragments).","commonSituations":"Reusing numpy-style concat code on WGMMA accumulator fragments; splitting accumulators along a non-leading dimension and trying to rejoin them; layout of intermediates changing to strided after a WGMMA op in a newer JAX.","solutions":["Concatenate along axis 0 instead (restructure the kernel so the joined dimension is the leading one).","Or go through memory: store fragments to a tiled/untiled reference, concatenate there, and reload with the strided layout.","Or convert fragments to a TiledLayout, concatenate on any axis, then convert back."],"exampleFix":"# before\nout = FragmentedArray.concatenate(accs, axis=1)  # WGStridedFragLayout -> error\n# after\n# transpose the split so joining happens along axis 0:\nout = FragmentedArray.concatenate(accs, axis=0)","handlingStrategy":"fallback","validationCode":"from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout\n\nif isinstance(arrays[0].layout, WGStridedFragLayout) and axis != 0:\n    axis = 0  # or restructure split so axis 0 is the joined dim\nout = FragmentedArray.concatenate(arrays, axis=axis)","typeGuard":"from jax.experimental.mosaic.gpu.fragmented_array import WGStridedFragLayout\n\ndef can_concat_axis0_only(arrays) -> bool:\n    return isinstance(arrays[0].layout, WGStridedFragLayout)","tryCatchPattern":"try:\n    out = FragmentedArray.concatenate(arrays, axis=axis)\nexcept NotImplementedError:\n    # fall back: concat along axis 0, or via TiledLayout round-trip\n    out = FragmentedArray.concatenate(arrays, axis=0)","preventionTips":["Keep the joined dimension as the leading axis for strided fragments.","Convert to TiledLayout when arbitrary-axis concat is needed.","Track which fragments come from WGMMA accumulators."],"tags":["jax","mosaic-gpu","not-implemented","wgmma","layout-limitation"],"backgroundTag":"unsupported-operation-not-implemented","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}