{"record":{"id":"c7371db74332b690","repo":"jax-ml/jax","slug":"arrays-with-the-splat-layout-can-only-be-stored-wh","errorCode":null,"errorMessage":"Arrays with the splat layout can only be stored when they have a single element or a multiple of {WARPGROUP_SIZE} elements","messagePattern":"Arrays with the splat layout can only be stored when they have a single element or a multiple of (.+?) elements","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3771,"sourceCode":"    ref = utils.memref_reshape(ref, (*(1 for _ in ref_ty.shape), *ref_ty.shape))\n    return cls.load_tiled(\n        ref, swizzle=swizzle, is_signed=is_signed, layout=layout, optimized=optimized\n    )\n\n  def _store_untiled_splat(self, ref: ir.Value):\n    if math.prod(self.shape) == 1:\n      c0 = c(0, ir.IndexType.get())\n      memref.store(\n          self.registers.flat[0], ref, [c0] * len(ir.MemRefType(ref.type).shape)\n      )\n      return\n\n    vec_size = 64 // mgpu.bitwidth(self.mlir_dtype)\n    if np.prod(self.shape) < vec_size * WARPGROUP_SIZE:\n      vec_size = 1\n\n    if np.prod(self.shape) % WARPGROUP_SIZE * vec_size:\n      raise NotImplementedError(\n          \"Arrays with the splat layout can only be stored when they have a\"\n          f\" single element or a multiple of {WARPGROUP_SIZE} elements\"\n      )\n\n    fa = FragmentedArray.splat(\n        self.registers.flat[0],\n        self.shape,\n        layout=WGStridedFragLayout(shape=self.shape, vec_size=vec_size),\n        is_signed=self.is_signed,\n    )\n    fa.store_untiled(ref)\n\n  def store_tiled_async(\n      self,\n      ref: ir.Value,\n      barrier: utils.BarrierRef,\n      cluster_dim: gpu.Dimension,\n      cluster_idx: ir.Value,","sourceCodeStart":3753,"sourceCodeEnd":3789,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3753-L3789","documentation":"Storing an array with the splat layout requires the total element count to be either a single element or an exact multiple of WARPGROUP_SIZE (128) — the hardware stores splat fragments in warpgroup-wide vector chunks of vec_size = 64/bitwidth elements, so other sizes cannot be lowered.","triggerScenarios":"FragmentedArray._store_untiled_splat / store on a WGSplatFragLayout array whose np.prod(shape) is not 1 and not a multiple of 128 (as combined with vec_size in the modulo check) — e.g. storing a splatted array of 3, 100, or 1000 elements.","commonSituations":"Splatting a scalar constant over an arbitrary-shaped result then storing it; block sizes not chosen as multiples of 128 (one warpgroup); changing tile shapes during kernel tuning so a splat operand's size is no longer warpgroup-aligned.","solutions":["Pick block/shape sizes that are multiples of WARPGROUP_SIZE (128) when a splat-layout array will be stored.","If the value is a true scalar, store a 1-element splat instead of splatting to the full shape.","Convert the array to a TiledLayout before storing if the shape can't be changed (fa.to_layout(TiledLayout(...))).store(ref))."],"exampleFix":"# before\nfa = FragmentedArray.splat(c)  # shape (100,)\nfa.store(ref)  # 100 not a multiple of 128\n\n# after\n# choose warpgroup-aligned block\nfa = FragmentedArray.splat(c)  # shape (128,)\nfa.store(ref)","handlingStrategy":"validation","validationCode":"import numpy as np\nfrom jax.experimental.mosaic.gpu.fragmented_array import WGSplatFragLayout, WARPGROUP_SIZE\nif isinstance(fa.layout, WGSplatFragLayout):\n    n = np.prod(fa.shape)\n    assert n == 1 or (n * (64 // mgpu.bitwidth(fa.mlir_dtype))) % WARPGROUP_SIZE == 0, 'splat store needs 1 or warpgroup-multiple elements'","typeGuard":"def splat_storeable(fa):\n    n = int(np.prod(fa.shape))\n    return n == 1 or n % WARPGROUP_SIZE == 0","tryCatchPattern":null,"preventionTips":["Choose block sizes as multiples of 128 (WARPGROUP_SIZE) when splat arrays are stored.","Store single-element splats for true scalars instead of splatting to full shapes.","Convert to TiledLayout via to_layout when the shape can't be aligned."],"tags":["jax","mosaic-gpu","splat-layout","warpgroup","alignment","store"],"backgroundTag":"memory-alignment-error","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}