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 Triton GPU backend because dot_general interprets all integer inputs as signed. Consider casting to a signed type before the dot operation.

What it means

The Pallas Triton GPU backend lowers lax.dot_general (matmul) to Triton's tl.dot, which interprets all integer inputs as signed. Therefore unsigned integer dtypes (uint8, uint16, uint32, uint64) in either operand are rejected at lowering time with NotImplementedError. Cast inputs to a signed dtype (e.g. int32) before the dot.

Source

Thrown at jax/_src/pallas/triton/lowering.py:2338

def _as_f32(x):
  return _ir_cast(x, _dtype_to_ir_type(jnp.float32), signed=False)


@register_lowering(lax.dot_general_p)
def _dot_general_lowering(
    ctx: LoweringRuleContext,
    a,
    b,
    *,
    dimension_numbers,
    out_sharding,
    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 Triton GPU backend because"
          " dot_general interprets all integer inputs as signed. Consider"
          " casting to a signed type before the dot operation."
      )
  del preferred_element_type, out_sharding  # Unused.
  ((a_contract_dim,), (b_contract_dim,)), batch_dims = dimension_numbers
  assert batch_dims == ((), ())

  if a_contract_dim == 0:
    a = tt_dialect.trans(a, (1, 0))
  if b_contract_dim == 1:
    b = tt_dialect.trans(b, (1, 0))

  a_aval, b_aval = ctx.avals_in
  [out_aval] = ctx.avals_out

  if precision is None or (precision == lax.DotAlgorithmPreset.DEFAULT):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast operands to a signed dtype before the dot: a = a.astype(jnp.int32), b = b.astype(jnp.int32)
2. If values fit, reinterpret with .view / lax.bitcast_convert_type to int8/int16 keeping bit pattern
3. Keep Pallas matmul kernels on float32/bfloat16/int8 inputs only

Example fix

# before
out = pl.dot(x_u8, w_u8)  # x_u8: uint8

# after
out = pl.dot(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32))

import jax.numpy as jnp

def ok_for_pallas_dot(*arrays):
    return all(not jnp.issubdtype(a.dtype, jnp.unsignedinteger) for a in arrays)

def is_signed_or_float(a) -> bool:
    import jax.numpy as jnp
    return not jnp.issubdtype(a.dtype, jnp.unsignedinteger)

When it happens

Trigger: Calling a Mosaic pallas (triton) kernel that performs pallas.tl.dot / jax.lax.dot_general on arrays whose aval dtype is any jnp.unsignedinteger subtype (e.g. uint8 block representations of weights).

Common situations: Loading quantized model weights as uint8 and feeding them directly to a Pallas matmul kernel; converting images stored as uint8 into a pallas kernel without normalization; porting a TPU pallas kernel that tolerated unsigned ints.

Related errors

• cannot cast from `{dtype_name}`
• cannot cast to `{dtype_name}`
• a and b must have the same element type, but got: {a_type.el
• {ctx.avals_out[0].dtype}
• Only float32 and int32 results are supported, got {dtype}

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