{"record":{"id":"4924cc2afeb737f3","repo":"jax-ml/jax","slug":"cannot-convert-from-self-layout-to-new-layout","errorCode":null,"errorMessage":"Cannot convert from {self.layout} to {new_layout}","messagePattern":"Cannot convert from (.+?) to (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/fragmented_array.py","lineNumber":1437,"sourceCode":"        #     reg[0]:   0 1 2 3 4 5 6 7       reg[2]:  16 17 18 19 20 21 22 23\n        #     prmt[0]:  -0- -1- -2- -3-                --4-- --5-- --6-- --7--\n        #     prmt[1]:  -4- -5- -6- -7-                --0-- --1-- --2-- --3--\n        # The expected outputs and their respective permutations are:\n        #     out[0]:   0 1 2 3 16 17 18 19   out[2]:  4 5 6 7 20 21 22 23\n        #     prmt[0]:  -0- -1- --4-- --5--  prmt[2]:  -6- -7- --2-- --3--\n        perm = arith.select(is_01, c(0x5410), c(0x3276))\n        blend = utils.prmt(reg, exchanged, perm)\n        for i in range(2):\n          reg = utils.vector_slice(blend, slice(i * 4, i * 4 + 4))\n          new_registers[(idx[0], idx[1] * 2 + i, *idx[2:-1])] = reg\n      assert all(r is not None for r in new_registers)\n      return FragmentedArray(\n          _registers=new_registers, _layout=new_layout, _is_signed=self.is_signed,\n      )\n    if self.layout == WGMMA_LAYOUT_UPCAST_4X and new_layout == WGMMA_LAYOUT:\n      return self.to_layout(WGMMA_LAYOUT_UPCAST_2X).to_layout(new_layout)\n    if not isinstance(self.layout, WGSplatFragLayout):\n      raise NotImplementedError(\n          f\"Cannot convert from {self.layout} to {new_layout}\"\n      )\n    return type(self).splat(\n        self.registers.item(), self.shape, new_layout, is_signed=self.is_signed\n    )\n\n  def _pointwise(\n      self,\n      op,\n      *other,\n      output_is_signed: bool | None = None,\n      restrict_bitwidth: bool = True,\n  ) -> FragmentedArray:\n    if restrict_bitwidth:\n      if (bitwidth := utils.bitwidth(self.mlir_dtype)) <= 8 and bitwidth != 1:\n        raise NotImplementedError(\n            f\"Pointwise operations on {bitwidth}-bit types are unsupported\"\n            \" (except bitwise operations). Upcast to a 16- or 32-bit type\"","sourceCodeStart":1419,"sourceCodeEnd":1455,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/fragmented_array.py#L1419-L1455","documentation":"to_layout handles conversions between known fragment layouts and, at the end of the chain, only knows how to finish by splatting when the source is a WGSplatFragLayout. If the source layout is a non-splat layout with no registered conversion to the target, the conversion is unimplemented and raises NotImplementedError.","triggerScenarios":"Calling to_layout(target) where no conversion path exists from self.layout to new_layout — e.g. converting a general distributed layout directly to a WGMMA layout, or to/from a layout pair never wired up in fragmented_array.py (the explicit WGSplatFragLayout and WGMMA_LAYOUT_UPCAST_4X→2X paths don't apply).","commonSituations":"Mixing layouts produced by different primitives (dot_general output layouts vs. memory layouts) in a custom Mosaic kernel; upgrading JAX where a previously implicit conversion was removed or restricted; writing new layout-cast ops against newly added layouts.","solutions":["Route through an intermediate layout that does have a conversion path (as the code itself does: UPCAST_4X → UPCAST_2X → WGMMA)","Materialize to registers/global memory and reload in the target layout instead of a direct cast","Check the set of supported conversions in fragmented_array.py and pick source/target layouts that are wired","Use FragmentedArray.splat if the value is uniform and the source is a splat layout"],"exampleFix":"# before\ny = x.to_layout(TARGET_LAYOUT)  # no direct path\n# after\ny = x.to_layout(INTERMEDIATE_LAYOUT).to_layout(TARGET_LAYOUT)","handlingStrategy":"fallback","validationCode":"SUPPORTED = {(type(a), type(b)) for known pairs}  # maintain from source\nassert (type(x.layout), type(target)) in SUPPORTED or isinstance(x.layout, WGSplatFragLayout)","typeGuard":null,"tryCatchPattern":"try:\n    return x.to_layout(target)\nexcept NotImplementedError:\n    for mid in INTERMEDIATES:\n        try:\n            return x.to_layout(mid).to_layout(target)\n        except NotImplementedError:\n            continue\n    raise","preventionTips":["Centralize layout conversions in one helper that tries known intermediate chains","Keep a table of tested layout pairs as regression tests"],"tags":["jax","mosaic-gpu","layout-conversion","not-implemented","wgmma"],"backgroundTag":"unsupported-layout-conversion","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}