{"record":{"id":"17533a35ce2e03db","repo":"jax-ml/jax","slug":"concatenating-arrays-with-splat-layout-is-not-supp","errorCode":null,"errorMessage":"Concatenating arrays with splat layout is not supported.","messagePattern":"Concatenating arrays with splat layout is not supported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":5414,"sourceCode":"        )\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":5396,"sourceCodeEnd":5420,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L5396-L5420","documentation":"FragmentedArray.concat handles strided-fragment layouts but explicitly does not support WGSplatFragLayout (a broadcast/splat register layout). The match statement raises NotImplementedError as a guard because concatenating splatted registers has no defined semantics in this code path. Developers must materialize or reshape the splat array first.","triggerScenarios":"Calling FragmentedArray.concat where any (here the first, arr0) array has WGSplatFragLayout, e.g. concatenating an array created from a scalar broadcast with normal strided-fragment arrays.","commonSituations":"Broadcasting a scalar constant into a kernel and then concatenating it with computed fragments; using older Mosaic versions where splat arrays silently worked differently after layout refactors.","solutions":["Materialize the splat array into a strided/regular fragment layout (e.g. broadcast its registers into full storage) before concat","Compute the concatenated result another way: pad or reshape operands so splat values are produced post-concatenation","File/upstream a feature request if concat-of-splat is needed, since this is NotImplementedError by design"],"exampleFix":"# before\nout = FragmentedArray.concat([splat_arr, other])  # splat_arr has WGSplatFragLayout\n\n# after\nmaterialized = splat_arr # convert to a strided fragment layout first\nout = FragmentedArray.concat([materialized, other])","handlingStrategy":"type-guard","validationCode":"from jax.experimental.mosaic.gpu import fragmented_array as fa\n\ndef concat_safe(arrays):\n  if any(isinstance(a.layout, fa.WGSplatFragLayout) for a in arrays):\n    raise TypeError('Splat-layout arrays cannot be concatenated; materialize them first')\n  return fa.FragmentedArray.concat(arrays)","typeGuard":"def is_splat(arr) -> bool:\n  from jax.experimental.mosaic.gpu import fragmented_array as fa\n  return isinstance(arr.layout, fa.WGSplatFragLayout)","tryCatchPattern":"try:\n  out = FragmentedArray.concat(arrays)\nexcept NotImplementedError as e:\n  if 'splat' in str(e):\n    arrays = [materialize(a) for a in arrays]  # your materialization\n    out = FragmentedArray.concat(arrays)\n  else:\n    raise","preventionTips":["Never pass broadcast/splat fragment arrays to concat","Check isinstance(arr.layout, WGSplatFragLayout) before concat","Keep splat constants out of concat operands; apply broadcasts after concatenation"],"tags":["mosaic-gpu","not-implemented","splat-layout","concat","fragmented-array"],"backgroundTag":"unsupported-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}