jax-ml/jax · error · ValueError

Only TiledLayout supports reductions.

Error message

Only TiledLayout supports reductions.

What it means

Raised when ReducedLayout.to_mgpu() resolves its inner layout to something other than mgpu.TiledLayout. Reductions (layout.reduce(axes)) are only defined for tiled layouts; reducing an elementwise/fragmented layout such as a transposed WGMMA layout is unsupported.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1839

    object.__setattr__(self, "kwargs", frozen_dict.FrozenDict(self.kwargs))

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    return self.layout_cls.to_mgpu(*self.args, **self.kwargs)


@dataclasses.dataclass(frozen=True)
class ReducedLayout(SomeLayout):
  layout: SomeLayout
  axes: Sequence[int]

  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:
    if args or kwargs:
      raise ValueError(f"Can't instantiate {self} with arguments.")
    layout = self.layout.to_mgpu()
    if not isinstance(layout, mgpu.TiledLayout):
      raise ValueError("Only TiledLayout supports reductions.")
    return layout.reduce(self.axes)


class Layout(SomeLayout, enum.Enum):
  #: [m, n] matrix, where m % 64 == 0 == n % 8.
  WGMMA = enum.auto()
  WGMMA_8BIT = enum.auto()
  WGMMA_UPCAST_2X = enum.auto()
  WGMMA_UPCAST_4X = enum.auto()
  WGMMA_TRANSPOSED = enum.auto()

  WG_SPLAT = enum.auto()
  WG_STRIDED = enum.auto()

  TILED = enum.auto()

  TCGEN05 = enum.auto()
  TCGEN05_TRANSPOSED = enum.auto()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a TiledLayout-compatible inner layout (e.g. a TILED/elementwise layout that resolves to mgpu.TiledLayout) before wrapping in ReducedLayout
2. Perform the reduction manually by reshaping the layout rather than using ReducedLayout

Example fix

// before
ReducedLayout(Layout.WGMMA_TRANSPOSED, axes=[1]).to_mgpu()

// after
ReducedLayout(Layout.TILED, axes=[1]).to_mgpu()  # inner resolves to TiledLayout

inner = reduced.layout.to_mgpu()
import jax._src.mosaic_gpu as mgpu
assert isinstance(inner, mgpu.TiledLayout), 'reductions need a TiledLayout inner'

def supports_reduction(reduced) -> bool:
    import jax._src.mosaic_gpu as mgpu
    return isinstance(reduced.layout.to_mgpu(), mgpu.TiledLayout)

try:
    return reduced.to_mgpu()
except ValueError as e:
    if 'TiledLayout' in str(e):
        # manual fallback: reshape/tile the layout yourself
        return manual_reduce(reduced.layout.to_mgpu(), reduced.axes)
    raise

When it happens

Trigger: Constructing ReducedLayout(layout=Layout.WGMMA_TRANSPOSED, axes=[1]) (inner to_mgpu() returns a non-TiledLayout) and then calling to_mgpu().

Common situations: Wrapping an arbitrary layout in ReducedLayout when reducing accumulator layouts for collective MMA; refactoring where the inner layout silently changed type.

Related errors

• packed cannot be specified if layout is specified.
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't instantiate {self} with arguments.
• Expected WGSplatFragLayout, but got {fa.layout}

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