jax-ml/jax · error · ValueError

Only signed accumulator supported for integer operands.

Error message

Only signed accumulator supported for integer operands.

What it means

Integer MMA requires a signed s32 accumulator (PTX wgmma/imma semantics); an unsigned or sign-unspecified accumulator raises ValueError.

Source

Thrown at jax/experimental/mosaic/gpu/mma.py:216

  # todo(cperivol): A tile shape can have dimensions that are higher
  # multiples of the mma op size as long as those dimensions are not
  # sharded across warps.
  i4 = ir.IntegerType.get_signless(4)
  i8 = ir.IntegerType.get_signless(8)
  i32 = ir.IntegerType.get_signless(32)
  bf16 = ir.BF16Type.get()
  f16 = ir.F16Type.get()
  f8e4m3fn = ir.Float8E4M3FNType.get()
  f8e5m2 = ir.Float8E5M2Type.get()
  if (element_type := a.mlir_dtype) != b.mlir_dtype:
    raise ValueError(f"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}")
  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):
    raise NotImplementedError(f"Unsupported operand type: {element_type}")
  if isinstance(element_type, ir.IntegerType):
    if acc.mlir_dtype != i32:
      raise NotImplementedError("Only s32 accumulator supported for integer operands.")
    if not acc.is_signed:
      raise ValueError("Only signed accumulator supported for integer operands.")
  elif acc.mlir_dtype != ir.F32Type.get():
    raise NotImplementedError("Only f32 accumulator supported for floating operands.")

  can_infer_from_acc_layout = (
      isinstance(acc.layout, fa.TiledLayout)
      and len(acc.layout.base_tile_shape) == 2
      and acc.layout.base_tile_shape[0] % 16 == 0
  )
  if not can_infer_from_acc_layout:
    raise ValueError("Expected MMALayouts.acc for acc")
  m_warps = acc.layout.base_tile_shape[0] // 16  # type: ignore
  layouts = MMALayouts(element_type, m_warps=m_warps)
  if layouts.lhs != a.layout:
    raise ValueError("Expected MMALayouts.lhs layout for A")
  if layouts.rhs != b.layout:
    raise ValueError("Expected MMALayouts.rhs layout for B")
  if layouts.acc != acc.layout:
    raise ValueError("Expected MMALayouts.acc layout for acc")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jnp.int32 (signed) when creating the accumulator
2. Pass explicit signedness metadata if constructing FragmentedArrays manually

Example fix

// before
acc = fa.from_tensor(np.zeros((m, n), np.uint32))
// after
acc = fa.from_tensor(jnp.zeros((m, n), jnp.int32))

if isinstance(a.mlir_dtype, ir.IntegerType):
    assert acc.is_signed, 'integer MMA needs signed s32 acc'

def is_signed_i32_acc(acc):
    return acc.is_signed and acc.mlir_dtype == ir.IntegerType.get_signless(32)

When it happens

Trigger: Calling mma() with integer operands and an accumulator FraggedArray whose is_signed is False/None, e.g. built from a signless ir.IntegerType(32).

Common situations: Creating the accumulator from raw MLIR types or unsigned numpy dtypes instead of jnp.int32.

Related errors

• is_signed must be specified for integer types
• Only s32 accumulator supported for integer operands.
• Acc ref must be at least 2D, got shape {shape}
• m_warps must be 1, 2, or 4, but got {m_warps=}
• M mismatch: {m} != {m2}

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