jax-ml/jax · error · NotImplementedError

Cannot convert from {self.layout} to {new_layout}

Error message

Cannot convert from {self.layout} to {new_layout}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/fragmented_array.py:1437

        #     reg[0]:   0 1 2 3 4 5 6 7       reg[2]:  16 17 18 19 20 21 22 23
        #     prmt[0]:  -0- -1- -2- -3-                --4-- --5-- --6-- --7--
        #     prmt[1]:  -4- -5- -6- -7-                --0-- --1-- --2-- --3--
        # The expected outputs and their respective permutations are:
        #     out[0]:   0 1 2 3 16 17 18 19   out[2]:  4 5 6 7 20 21 22 23
        #     prmt[0]:  -0- -1- --4-- --5--  prmt[2]:  -6- -7- --2-- --3--
        perm = arith.select(is_01, c(0x5410), c(0x3276))
        blend = utils.prmt(reg, exchanged, perm)
        for i in range(2):
          reg = utils.vector_slice(blend, slice(i * 4, i * 4 + 4))
          new_registers[(idx[0], idx[1] * 2 + i, *idx[2:-1])] = reg
      assert all(r is not None for r in new_registers)
      return FragmentedArray(
          _registers=new_registers, _layout=new_layout, _is_signed=self.is_signed,
      )
    if self.layout == WGMMA_LAYOUT_UPCAST_4X and new_layout == WGMMA_LAYOUT:
      return self.to_layout(WGMMA_LAYOUT_UPCAST_2X).to_layout(new_layout)
    if not isinstance(self.layout, WGSplatFragLayout):
      raise NotImplementedError(
          f"Cannot convert from {self.layout} to {new_layout}"
      )
    return type(self).splat(
        self.registers.item(), self.shape, new_layout, is_signed=self.is_signed
    )

  def _pointwise(
      self,
      op,
      *other,
      output_is_signed: bool | None = None,
      restrict_bitwidth: bool = True,
  ) -> FragmentedArray:
    if restrict_bitwidth:
      if (bitwidth := utils.bitwidth(self.mlir_dtype)) <= 8 and bitwidth != 1:
        raise NotImplementedError(
            f"Pointwise operations on {bitwidth}-bit types are unsupported"
            " (except bitwise operations). Upcast to a 16- or 32-bit type"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Route through an intermediate layout that does have a conversion path (as the code itself does: UPCAST_4X → UPCAST_2X → WGMMA)
2. Materialize to registers/global memory and reload in the target layout instead of a direct cast
3. Check the set of supported conversions in fragmented_array.py and pick source/target layouts that are wired
4. Use FragmentedArray.splat if the value is uniform and the source is a splat layout

Example fix

# before
y = x.to_layout(TARGET_LAYOUT)  # no direct path
# after
y = x.to_layout(INTERMEDIATE_LAYOUT).to_layout(TARGET_LAYOUT)

SUPPORTED = {(type(a), type(b)) for known pairs}  # maintain from source
assert (type(x.layout), type(target)) in SUPPORTED or isinstance(x.layout, WGSplatFragLayout)

try:
    return x.to_layout(target)
except NotImplementedError:
    for mid in INTERMEDIATES:
        try:
            return x.to_layout(mid).to_layout(target)
        except NotImplementedError:
            continue
    raise

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• Non-trivial indexing on WGMMAAbstractAccumulatorRef is not s
• Concatenating arrays with strided layout is only supported a
• Unsupported dtype: {ref.dtype}
• Only SMEM and TMEM refs are supported.
• Unsupported transform: {type(transform)}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/4924cc2afeb737f3. Report an issue: GitHub.