jax-ml/jax · error · ValueError

is_signed must be specified for integer types

Error message

is_signed must be specified for integer types

What it means

_ptx_dtype_str builds the PTX dtype string for MMA instructions. Integer MLIR types carry no sign, so the caller must pass is_signed to choose the 's'/'u' prefix; if None for an integer dtype, ValueError.

Source

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

        _check_canonical=False,
    ).canonicalize()
    self.acc = fa.TiledLayout(
        fa.Tiling(((m_warps * 16, n_warps * 8), (16, 8), (8, 8), (2,))),
        warp_dims=(-7, -6),
        lane_dims=(-3, -2),
        vector_dim=-1,
        _check_canonical=False,
    ).canonicalize()


def _ptx_dtype_str(dtype: ir.Type, *, is_signed: bool | None = None) -> str:
  if isinstance(dtype, ir.Float8E4M3FNType):
    return "e4m3"
  elif isinstance(dtype, ir.Float8E5M2Type):
    return "e5m2"
  elif isinstance(dtype, ir.IntegerType):
    if is_signed is None:
      raise ValueError("is_signed must be specified for integer types")
    prefix = "s" if is_signed else "u"
    return f"{prefix}{dtype.width}"
  return str(dtype)


def _mma_single_tile(
    acc: fa.FragmentedArray, a: fa.FragmentedArray, b: fa.FragmentedArray
) -> fa.FragmentedArray:
  """Performs `acc + a @ b` using warp level MMA instructions."""
  i32 = ir.IntegerType.get_signless(32)

  k_tile = 256 // utils.bitwidth(a.mlir_dtype)
  assert a.mlir_dtype == b.mlir_dtype
  is_integer = isinstance(a.mlir_dtype, ir.IntegerType)
  assert acc.mlir_dtype == i32 if is_integer else ir.F32Type.get()
  assert acc.is_signed in {None, True}
  assert (
      isinstance(acc.layout, fa.TiledLayout)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Create integer FraggedArrays with explicit signedness (e.g. jnp int8 arrays convert with is_signed=True)
2. Pass is_signed=True/False wherever the helper accepts it
3. Use dtypes with inherent signedness (jnp.int8) instead of raw ir.IntegerType

Example fix

// before
a = fa.FragmentedArray(..., mlir_dtype=ir.IntegerType.get_signless(8))
// after
a = fa.FragmentedArray(..., mlir_dtype=dtypes.int8)  # signedness known

if isinstance(dtype, ir.IntegerType):
    assert is_signed is not None, 'integer MMA requires is_signed'

def has_known_sign(fa):
    return not isinstance(fa.mlir_dtype, ir.IntegerType) or fa.is_signed is not None

When it happens

Trigger: Calling mma/related helpers with i8 or i4 operands where the sign information wasn't provided (e.g. FragmentedArray built from plain ir.IntegerType without is_signed).

Common situations: Using int8/int4 MMA on Hopper with arrays created via arithmetic constants or from tensors lacking signedness metadata.

Related errors

• Only signed accumulator supported for integer operands.
• Sparse MMA unsupported for f32
• MMA with element type {elem_type_str} does not support block
• MMA with element type {elem_type_str} only supports accumula
• MMA with element type {elem_type_str} only supports accumula

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