{"record":{"id":"490d586f581090a5","repo":"jax-ml/jax","slug":"incompatible-fragmentedarray-layouts","errorCode":null,"errorMessage":"Incompatible FragmentedArray layouts","messagePattern":"Incompatible FragmentedArray layouts","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1505,"sourceCode":"          raise NotImplementedError(o)\n\n        o = FragmentedArray.splat(\n            o, shape=self.shape, layout=self.layout, is_signed=self.is_signed\n        )\n\n      if isinstance(o.layout, WGSplatFragLayout):\n        if not o.layout.can_broadcast_to(self.shape):\n          raise ValueError(\n              f\"Cannot broadcast shape {self.shape} to layout {o.layout}\")\n        o = FragmentedArray.splat(\n            o.registers.flat[0],\n            shape=self.shape,\n            layout=self.layout,\n            is_signed=o.is_signed,\n        )\n      else:\n        if self.layout != o.layout:\n          raise ValueError(\"Incompatible FragmentedArray layouts\")\n        if self.registers.shape != o.registers.shape:\n          raise ValueError(\"Incompatible FragmentedArray shapes\")\n\n      other_arrs.append(o)\n    new_regs = np.empty_like(self.registers)\n\n    for idx, reg in np.ndenumerate(self.registers):\n      new_regs[idx] = op(reg, *(o.registers[idx] for o in other_arrs))\n    reg_ty = new_regs.flat[0].type\n    if isinstance(reg_ty, ir.VectorType):\n      reg_ty = ir.VectorType(reg_ty).element_type\n    if output_is_signed is None and isinstance(reg_ty, ir.IntegerType):\n      output_is_signed = self.is_signed\n    return FragmentedArray(\n        _registers=new_regs, _layout=self.layout, _is_signed=output_is_signed\n    )\n\n  def __pos__(self):","sourceCodeStart":1487,"sourceCodeEnd":1523,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1487-L1523","documentation":"_pointwise requires both operands to end up in the identical fragment layout so registers can be paired elementwise. After splat handling, if self.layout != o.layout for a non-splat operand, there is no implicit conversion and the op aborts with ValueError.","triggerScenarios":"Applying +, -, *, neg, clz etc. between two FragmentedArrays with different layouts, e.g. a WGMMA-layout accumulator plus a memory/distributed-layout operand, without an explicit to_layout.","commonSituations":"Feeding a value loaded in one layout into arithmetic with a dot-product result in WGMMA layout; mixing arrays from different primitives in custom Mosaic kernels; layouts differing only in subtle parameters (num_regs, row/col splits) so the mismatch is invisible in logs.","solutions":["Explicitly convert one operand: other = other.to_layout(self.layout) before the op","Produce both operands with the same layout at creation time (same load/layout parameters)","Print/assert self.layout == other.layout in kernel debug builds to catch drift early"],"exampleFix":"# before\nz = acc + rhs  # acc.layout != rhs.layout\n# after\nz = acc + rhs.to_layout(acc.layout)","handlingStrategy":"validation","validationCode":"assert self.layout == o.layout, f\"layout mismatch: {self.layout} vs {o.layout}\"","typeGuard":"def same_layout(a, b) -> bool:\n    return a.layout == b.layout","tryCatchPattern":"try:\n    z = a + b\nexcept ValueError as e:\n    if 'Incompatible FragmentedArray layouts' in str(e):\n        z = a + b.to_layout(a.layout)\n    else:\n        raise","preventionTips":["Convert operands to a canonical layout at kernel entry","Add layout equality asserts in debug builds of custom primitives"],"tags":["jax","mosaic-gpu","layout-mismatch","pointwise","fragmented-array"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}