{"record":{"id":"c5656f4cb156eb78","repo":"jax-ml/jax","slug":"incompatible-fragmentedarray-shapes","errorCode":null,"errorMessage":"Incompatible FragmentedArray shapes","messagePattern":"Incompatible FragmentedArray shapes","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1507,"sourceCode":"        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):\n    return self\n","sourceCodeStart":1489,"sourceCodeEnd":1525,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1489-L1525","documentation":"Even with identical layouts, two FragmentedArrays can hold different numbers of register tiles (registers_shape depends on shape and layout parameters). _pointwise pairs registers index-by-index, so unequal register arrays cannot be combined and raise ValueError.","triggerScenarios":"Applying arithmetic between arrays with the same layout but different element shapes — e.g. (128, 8) + (128, 16) in the same layout — so registers.shape differs; also shapes that are broadcast-compatible in NumPy but not in register-space Mosaic.","commonSituations":"Assuming NumPy broadcasting semantics for Mosaic fragments; padding one operand's tile shape differently; partial-tile kernels where one operand was padded to a full tile and the other was not.","solutions":["Reshape/pad both operands to the identical shape before the pointwise op","If broadcasting was intended, splat the smaller operand to the larger shape (FragmentedArray.splat / broadcast) instead of relying on NumPy rules","Assert o.registers.shape == self.registers.shape in tests to catch silent tile-shape divergence"],"exampleFix":"# before\nz = a + b  # a.shape=(128,8), b.shape=(128,16)\n# after\nb = b[:, :8]  # or pad a to (128,16)\nz = a + b","handlingStrategy":"validation","validationCode":"assert self.registers.shape == o.registers.shape, (\n    f\"register shapes differ: {self.registers.shape} vs {o.registers.shape}\")","typeGuard":"def same_registers(a, b) -> bool:\n    return a.registers.shape == b.registers.shape","tryCatchPattern":"try:\n    z = a + b\nexcept ValueError as e:\n    if 'Incompatible FragmentedArray shapes' in str(e):\n        raise ValueError(f\"pad/reshape operands {a.shape} and {b.shape} to match\") from e\n    raise","preventionTips":["Don't assume NumPy broadcasting for FragmentedArray operands","Pad partial tiles consistently for every operand in a fused expression"],"tags":["jax","mosaic-gpu","shape-mismatch","registers","pointwise"],"backgroundTag":"operand-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}