{"record":{"id":"49417229fc604ce2","repo":"jax-ml/jax","slug":"unsigned-integer-dtype-aval-dtype-is-not-support","errorCode":null,"errorMessage":"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.","messagePattern":"Unsigned integer 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\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2839,"sourceCode":"      f\"#tpu.dot_dimension_numbers<{','.join(map(format_dims, all_dims))}>\"\n  )\n\n  return ir.Attribute.parse(tpu_dim_numbers_str)\n\n\n@register_lowering_rule(lax.dot_general_p)\ndef _dot_general_lowering_rule(\n    ctx: LoweringRuleContext,\n    x,\n    y,\n    dimension_numbers,\n    precision,\n    preferred_element_type,\n    **_,\n):\n  for aval in ctx.avals_in:\n    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):\n      raise NotImplementedError(\n          f\"Unsigned integer dtype {aval.dtype} is not supported for\"\n          \" dot_general (matmul) on the Pallas Mosaic TPU backend because\"\n          \" dot_general interprets all integer inputs as signed. Consider\"\n          \" casting to a signed type before the dot operation.\"\n      )\n  (lhs_dims, rhs_dims), _ = dimension_numbers\n  (aval_out,) = ctx.avals_out\n  out_type = ctx.aval_to_ir_type(aval_out)\n  assert isinstance(out_type, ir.ShapedType)\n  val_type = ir.ShapedType(out_type).element_type\n  if any(\n      isinstance(val_type, cls)\n      for cls in [\n          ir.BF16Type,\n          ir.F32Type,\n          ir.Float8E5M2Type,\n          ir.Float8E4M3FNType,\n          ir.Float8E4M3B11FNUZType,","sourceCodeStart":2821,"sourceCodeEnd":2857,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2821-L2857","documentation":"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.","triggerScenarios":"Calling jax.numpy.matmul / lax.dot_general with uint8/uint16/uint32 operands inside a Pallas Mosaic TPU kernel.","commonSituations":"Kernels over quantized or image data stored as uint8; reusing GPU quantization code on TPU.","solutions":["Cast operands to a signed dtype before the dot: x.astype(jnp.int32), then adjust result interpretation","Store/pre-quantize data as int8 instead of uint8 before the kernel","Keep dot operands in float/bf16 if unsigned semantics aren't essential"],"exampleFix":"// before\nout = jnp.matmul(x_u8, w_u8)\n// after\nout = jnp.matmul(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\ndef dot_safe(a, b):\n    if jnp.issubdtype(a.dtype, jnp.unsignedinteger) or jnp.issubdtype(b.dtype, jnp.unsignedinteger):\n        a = a.astype(jnp.int32); b = b.astype(jnp.int32)\n    return jnp.matmul(a, b)","typeGuard":"def is_signed_for_dot(dt) -> bool:\n    import jax.numpy as jnp\n    return not jnp.issubdtype(dt, jnp.unsignedinteger)","tryCatchPattern":null,"preventionTips":["Validate kernel input dtypes at the boundary before launching","Standardize on signed dtypes for TPU kernels","Add dtype asserts in kernel wrappers"],"tags":["jax","pallas","tpu","dot-general","unsigned-integer","dtype"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}