jax-ml/jax · error · NotImplementedError

Unsigned integer dtype {aval.dtype} is not supported for dot

Error message

Unsigned integer dtype {aval.dtype} is not supported for dot_general (matmul) on the Pallas Mosaic TPU backend because dot_general interprets all integer inputs as signed. Consider casting to a signed type before the dot operation.

What it means

Mosaic TPU's dot_general lowering treats all integer inputs as signed, so unsigned integer dtypes are explicitly rejected. The hardware contract has no unsigned interpretation, so the check fails fast with a suggestion to cast.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:2839

      f"#tpu.dot_dimension_numbers<{','.join(map(format_dims, all_dims))}>"
  )

  return ir.Attribute.parse(tpu_dim_numbers_str)


@register_lowering_rule(lax.dot_general_p)
def _dot_general_lowering_rule(
    ctx: LoweringRuleContext,
    x,
    y,
    dimension_numbers,
    precision,
    preferred_element_type,
    **_,
):
  for aval in ctx.avals_in:
    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):
      raise NotImplementedError(
          f"Unsigned integer dtype {aval.dtype} is not supported for"
          " dot_general (matmul) on the Pallas Mosaic TPU backend because"
          " dot_general interprets all integer inputs as signed. Consider"
          " casting to a signed type before the dot operation."
      )
  (lhs_dims, rhs_dims), _ = dimension_numbers
  (aval_out,) = ctx.avals_out
  out_type = ctx.aval_to_ir_type(aval_out)
  assert isinstance(out_type, ir.ShapedType)
  val_type = ir.ShapedType(out_type).element_type
  if any(
      isinstance(val_type, cls)
      for cls in [
          ir.BF16Type,
          ir.F32Type,
          ir.Float8E5M2Type,
          ir.Float8E4M3FNType,
          ir.Float8E4M3B11FNUZType,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast operands to a signed dtype before the dot: x.astype(jnp.int32), then adjust result interpretation
2. Store/pre-quantize data as int8 instead of uint8 before the kernel
3. Keep dot operands in float/bf16 if unsigned semantics aren't essential

Example fix

// before
out = jnp.matmul(x_u8, w_u8)
// after
out = jnp.matmul(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32))

import jax.numpy as jnp
def dot_safe(a, b):
    if jnp.issubdtype(a.dtype, jnp.unsignedinteger) or jnp.issubdtype(b.dtype, jnp.unsignedinteger):
        a = a.astype(jnp.int32); b = b.astype(jnp.int32)
    return jnp.matmul(a, b)

def is_signed_for_dot(dt) -> bool:
    import jax.numpy as jnp
    return not jnp.issubdtype(dt, jnp.unsignedinteger)

When it happens

Trigger: Calling jax.numpy.matmul / lax.dot_general with uint8/uint16/uint32 operands inside a Pallas Mosaic TPU kernel.

Common situations: Kernels over quantized or image data stored as uint8; reusing GPU quantization code on TPU.

Related errors

• Reductions over unsigned integers not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for con
• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over {x_aval.dtype} not implemented.

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