{"record":{"id":"30252df5cc30698f","repo":"jax-ml/jax","slug":"source-and-destination-layouts-aren-t-compatible-f","errorCode":null,"errorMessage":"Source and destination layouts aren't compatible for a broadcast","messagePattern":"Source and destination layouts aren't compatible for a broadcast","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":3512,"sourceCode":"      return FragmentedArray(\n          _registers=np.tile(\n              self.registers,\n              math.prod(shape) // math.prod(self.shape),\n          ),\n          _layout=layout,\n          _is_signed=self.is_signed,\n      )\n    if not isinstance(self.layout, TiledLayout) or not isinstance(layout, TiledLayout):\n      raise NotImplementedError(self.layout, layout)\n    if len(layout.base_tile_shape) != len(shape):\n      raise NotImplementedError(\n          \"Tiling rank different than broadcast result rank, \"\n          f\"{layout.base_tile_shape} vs {shape}\"\n      )\n    new_dimensions = sorted(set(range(len(shape))) - set(source_dimensions))\n    expected_layout = layout.reduce(new_dimensions)\n    if expected_layout != self.layout:\n      raise ValueError(\n          \"Source and destination layouts aren't compatible for a broadcast\"\n      )\n    new_registers_shape = layout.registers_shape(shape)\n    pre_broadcast_registers_shape = list(new_registers_shape)\n    for new_dim in new_dimensions:\n      for i, is_new in enumerate(layout.tiling.tile_dimension(new_dim)):\n        if is_new:\n          pre_broadcast_registers_shape[i] = 1\n    # The broadcast for all dims but the vector_dim amounts to repeating the\n    # registers along the new dimensions. Along the vector_dim, we actually need\n    # to extend the vector length to change the type of the registers.\n    if layout.vector_length != self.layout.vector_length:\n      assert self.layout.vector_length == 1\n      registers = np.empty_like(self.registers)\n      for idx, reg in np.ndenumerate(self.registers):\n        registers[idx] = utils.vector_concat([reg] * layout.vector_length)\n    else:\n      registers = self.registers","sourceCodeStart":3494,"sourceCodeEnd":3530,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L3494-L3530","documentation":"Thrown by FragmentedArray broadcast when the destination layout cannot be produced by reducing (removing) the newly added broadcast dimensions from the source layout. Mosaic GPU layouts must be structurally compatible across a broadcast: the non-broadcast dims of the result layout must exactly match the source layout with the new dims collapsed. If they don't match, broadcasting register fragments between layouts is not supported.","triggerScenarios":"Calling FragmentedArray.broadcast (or ops that lower to it, e.g. _pointwise with shape promotion, to_layout, __getitem__, custom primitive blocks) where target_shape adds dimensions whose reduced layout differs from self.layout — e.g. mixing WGStridedFragLayout/WGSplatFragLayout with a tiled result, or broadcasting to a shape whose tiling/vec dims don't line up with the source's.","commonSituations":"Writing Mosaic GPU kernels where an operand with a warpgroup fragment layout is broadcast against an array with a TiledLayout; manual layout construction via to_layout before a binary op; shape changes that add dims incompatible with the chosen layout.","solutions":["Check the two layouts: assert layout.reduce(new_dims) == self.layout before broadcasting; align operand layouts with fg.to_layout or mgpu.repeat/lower operands to a common TiledLayout first.","Avoid broadcasting fragment-layout arrays directly; instead splat the scalar/operand explicitly (FragmentedArray.splat) or re-fragment the operand so both sides share one layout.","If the layouts look compatible, verify target_shape rank and that only new dims are being added — a mismatched existing dim also fails the equality check."],"exampleFix":"# before\nresult = a.broadcast((m, n))  # a has WGStridedFragLayout, fails\n\n# after\nb = mgpu.fragmented_array(..., layout=TiledLayout(...))\na2 = a.to_layout(b.layout)\nresult = a2.broadcast((m, n))","handlingStrategy":"validation","validationCode":"new_dims = sorted(set(range(len(target_shape))) - set(range(len(a.shape))))\nassert a.layout.reduce(new_dims) == a.layout_expected, 'broadcast layouts incompatible'","typeGuard":"def can_broadcast(a, target_shape, src_dims):\n    new_dims = sorted(set(range(len(target_shape))) - set(src_dims))\n    try:\n        return a.layout.reduce(new_dims) == a.layout\n    except Exception:\n        return False","tryCatchPattern":"try:\n    out = a.broadcast(shape)\nexcept ValueError as e:\n    if 'compatible for a broadcast' in str(e):\n        a = a.to_layout(common_layout)\n        out = a.broadcast(shape)\n    else:\n        raise","preventionTips":["Keep both operands of pointwise ops on the same layout (convert with to_layout before combining).","Prefer TiledLayout for arrays that will be broadcast.","Unit-test kernel wrappers with all intended shapes to catch layout mismatches early."],"tags":["jax","mosaic-gpu","layout","broadcast","fragmented-array"],"backgroundTag":"layout-incompatibility","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}