{"record":{"id":"8fe88fa7ba58c288","repo":"jax-ml/jax","slug":"only-signed-accumulator-supported-for-integer-oper","errorCode":null,"errorMessage":"Only signed accumulator supported for integer operands.","messagePattern":"Only signed accumulator supported for integer operands\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/mma.py","lineNumber":216,"sourceCode":"  # todo(cperivol): A tile shape can have dimensions that are higher\n  # multiples of the mma op size as long as those dimensions are not\n  # sharded across warps.\n  i4 = ir.IntegerType.get_signless(4)\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  bf16 = ir.BF16Type.get()\n  f16 = ir.F16Type.get()\n  f8e4m3fn = ir.Float8E4M3FNType.get()\n  f8e5m2 = ir.Float8E5M2Type.get()\n  if (element_type := a.mlir_dtype) != b.mlir_dtype:\n    raise ValueError(f\"Dtype mismatch: {a.mlir_dtype} != {b.mlir_dtype}\")\n  if element_type not in (bf16, f16, f8e4m3fn, f8e5m2, i8, i4):\n    raise NotImplementedError(f\"Unsupported operand type: {element_type}\")\n  if isinstance(element_type, ir.IntegerType):\n    if acc.mlir_dtype != i32:\n      raise NotImplementedError(\"Only s32 accumulator supported for integer operands.\")\n    if not acc.is_signed:\n      raise ValueError(\"Only signed accumulator supported for integer operands.\")\n  elif acc.mlir_dtype != ir.F32Type.get():\n    raise NotImplementedError(\"Only f32 accumulator supported for floating operands.\")\n\n  can_infer_from_acc_layout = (\n      isinstance(acc.layout, fa.TiledLayout)\n      and len(acc.layout.base_tile_shape) == 2\n      and acc.layout.base_tile_shape[0] % 16 == 0\n  )\n  if not can_infer_from_acc_layout:\n    raise ValueError(\"Expected MMALayouts.acc for acc\")\n  m_warps = acc.layout.base_tile_shape[0] // 16  # type: ignore\n  layouts = MMALayouts(element_type, m_warps=m_warps)\n  if layouts.lhs != a.layout:\n    raise ValueError(\"Expected MMALayouts.lhs layout for A\")\n  if layouts.rhs != b.layout:\n    raise ValueError(\"Expected MMALayouts.rhs layout for B\")\n  if layouts.acc != acc.layout:\n    raise ValueError(\"Expected MMALayouts.acc layout for acc\")","sourceCodeStart":198,"sourceCodeEnd":234,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/mma.py#L198-L234","documentation":"Integer MMA requires a signed s32 accumulator (PTX wgmma/imma semantics); an unsigned or sign-unspecified accumulator raises ValueError.","triggerScenarios":"Calling mma() with integer operands and an accumulator FraggedArray whose is_signed is False/None, e.g. built from a signless ir.IntegerType(32).","commonSituations":"Creating the accumulator from raw MLIR types or unsigned numpy dtypes instead of jnp.int32.","solutions":["Use jnp.int32 (signed) when creating the accumulator","Pass explicit signedness metadata if constructing FragmentedArrays manually"],"exampleFix":"// before\nacc = fa.from_tensor(np.zeros((m, n), np.uint32))\n// after\nacc = fa.from_tensor(jnp.zeros((m, n), jnp.int32))","handlingStrategy":"validation","validationCode":"if isinstance(a.mlir_dtype, ir.IntegerType):\n    assert acc.is_signed, 'integer MMA needs signed s32 acc'","typeGuard":"def is_signed_i32_acc(acc):\n    return acc.is_signed and acc.mlir_dtype == ir.IntegerType.get_signless(32)","tryCatchPattern":null,"preventionTips":["Use jnp.int32 zeros for accumulators in quantized kernels"],"tags":["jax","mosaic","mma","accumulator","signedness"],"backgroundTag":"accumulator-dtype-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}