jax-ml/jax · error · NotImplementedError

Only 1 reduction dimension is supported.

Error message

Only 1 reduction dimension is supported.

What it means

Mosaic only implements multi_dim_reduction over exactly one dimension; reducing 2+ dims at once is unsupported.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:889

      raise NotImplementedError(f"Unsupported reduction kind: {op.kind}")
  assert isinstance(result.layout, fa.WGSplatFragLayout)
  return [result.registers.item()]


@_register_lowering(vector.MultiDimReductionOp)
def _vector_multi_dim_reduction_op_lowering_rule(
    ctx: LoweringContext, op: vector.MultiDimReductionOp
) -> Sequence[ir.Value]:
  [in_layout, acc_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  if out_layout != acc_layout:
    raise ValueError(
        f"Output layout {out_layout} must match the accumulator layout"
        f" {acc_layout}"
    )

  if len(op.reduction_dims) != 1:
    raise NotImplementedError("Only 1 reduction dimension is supported.")

  op_kind = _combining_kind(op.kind)
  is_signed = _is_reduction_signed(op_kind)
  src = _fragmented_array_from_ir(op.source, in_layout, is_signed)
  acc = _fragmented_array_from_ir(op.acc, acc_layout, is_signed)

  if not isinstance(src.layout, fa.TiledLayout):
    raise NotImplementedError(f"Unsupported layout: {src.layout}")
  reduced_dim = src.layout.tiling.tile_dimension(op.reduction_dims[0])
  if any(reduced_dim[d] for d in src.layout.partitioned_warp_dims):
    # cross-warp reductions require scratch space.
    dtype = op.source.type.element_type
    allocation_size = ir.IntegerAttr(op.attributes["scratch_size"]).value * 8 // utils.bitwidth(dtype)
    scratch = _slice_smem(
        ir.MemRefType.get([allocation_size], dtype, memory_space=utils.smem()),
        ir.IntegerAttr(op.attributes["offset"]).value,
        ctx.smem_requested_bytes,
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Split into successive single-dimension reductions
2. Reduce over the tiled dimension last after reshaping to 2D

Example fix

// before
r = reduction(v, acc, axes=(1, 2))
// after
r = reduction(v, acc, axes=(1,))
r = reduction(r, acc0, axes=(1,))

assert len(op.reduction_dims) == 1, 'reduce one dim at a time'

When it happens

Trigger: vector.multi_dim_reduction with len(op.reduction_dims) != 1.

Common situations: Reducing a 3D tile over two axes in one op instead of chaining per-axis reductions.

Understand the failure class

Background: UnsupportedOperationException and "is not supported" errors: when a library deliberately refuses a call — this error's family across 30 libraries.

Related errors

• Unsupported reduction kind: {op.kind}
• 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}

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