jax-ml/jax · error · ValueError

Fused load-reduce is not supported for this layout

Error message

Fused load-reduce is not supported for this layout

What it means

When loading from a default TMEM layout (>=16-bit elements) directly into the LAYOUT register layout via _load_32xcols, no fused reduction is performed, so passing a reduce argument is rejected.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1372

    has_default_layout = self.layout == tmem_default_layout(packing)
    regs_shape = layout.registers_shape(self.shape)
    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the
    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.
    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to
    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.
    # This restriction can be lifted if `TiledLayout` supports multiple
    # vector dims.
    if self.layout == sparse_meta_layout():
      raise NotImplementedError("Sparse meta layout loads unsupported.")
    if regs_shape[0] != 1:  # We'll need to issue multiple loads below.
      raise NotImplementedError("Loading multiple row tiles")
    if (
        layout == LAYOUT
        and self.layout == tmem_default_layout(packing)
        and is_at_least_16b
    ):
      if reduce is not None:
        raise ValueError(
            "Fused load-reduce is not supported for this layout"
        )
      reduced_reg = None
      registers = _load_32xcols(
          self.address, columns, self.dtype, packing
      ).T.reshape(regs_shape)
    elif layout == self.layout.as_tiled_layout() and packing * bitwidth == 32:
      # TODO(apaszke): We raise NotImplemented here because technically for some
      # layouts this does make sense. I think only for those where all
      # dimensions that map to columns map only to columns.
      if reduce is not None:
        raise NotImplementedError(
            "Fused load-reduce is not supported for this layout"
        )
      assert len(layout.base_tile_shape) == 2
      cols = math.prod(regs_shape) * packing
      flat_registers, reduced_reg = _load_32xcols_native(
          self.address, cols, self.dtype, packing, packing, reduce=None

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Call load without reduce and reduce in registers afterwards
2. Use a TMEM/register layout pair that supports fused load-reduce (the _load_32xcols_native path)
3. Only pass reduce when the layout conditions permit it

Example fix

// before
arr = tmem.load(tcgen05.LAYOUT, reduce='max')
// after
arr = tmem.load(tcgen05.LAYOUT)
red = arr.reduce('max', axis=0)

if reduce is not None and tmem.layout == tcgen05.tmem_default_layout(tmem.packing):
    reduce = None  # fused reduce unsupported on this path

try:
    arr, red = tmem.load(layout, reduce=reduce)
except ValueError:
    arr, _ = tmem.load(layout); red = manual_reduce(arr)

When it happens

Trigger: tmem.load(LAYOUT, reduce=...) where self.layout == tmem_default_layout(packing) and element bitwidth is 16 or 32.

Common situations: Writing generic load-reduce helper code that always passes reduce; switching from a layout that supports fused reduce (e.g. native tiled) to the default TMEM layout.

Related errors

• Output layout {out_layout} must match the accumulator layout
• Unsupported layout: {src.layout}
• Unsupported reduction for f32. Only min, max, absmin, and ab
• Unsupported dtype for reduction: {self.dtype}
• Loading multiple row tiles

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