{"record":{"id":"e557e7c4797346c2","repo":"jax-ml/jax","slug":"unsigned-integer-dtype-aval-dtype-is-not-support-e557e7","errorCode":null,"errorMessage":"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.","messagePattern":"Unsigned integer 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\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/triton/lowering.py","lineNumber":2338,"sourceCode":"\ndef _as_f32(x):\n  return _ir_cast(x, _dtype_to_ir_type(jnp.float32), signed=False)\n\n\n@register_lowering(lax.dot_general_p)\ndef _dot_general_lowering(\n    ctx: LoweringRuleContext,\n    a,\n    b,\n    *,\n    dimension_numbers,\n    out_sharding,\n    precision,\n    preferred_element_type,\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 Triton GPU backend because\"\n          \" dot_general interprets all integer inputs as signed. Consider\"\n          \" casting to a signed type before the dot operation.\"\n      )\n  del preferred_element_type, out_sharding  # Unused.\n  ((a_contract_dim,), (b_contract_dim,)), batch_dims = dimension_numbers\n  assert batch_dims == ((), ())\n\n  if a_contract_dim == 0:\n    a = tt_dialect.trans(a, (1, 0))\n  if b_contract_dim == 1:\n    b = tt_dialect.trans(b, (1, 0))\n\n  a_aval, b_aval = ctx.avals_in\n  [out_aval] = ctx.avals_out\n\n  if precision is None or (precision == lax.DotAlgorithmPreset.DEFAULT):","sourceCodeStart":2320,"sourceCodeEnd":2356,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/triton/lowering.py#L2320-L2356","documentation":"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.","triggerScenarios":"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).","commonSituations":"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.","solutions":["Cast operands to a signed dtype before the dot: a = a.astype(jnp.int32), b = b.astype(jnp.int32)","If values fit, reinterpret with .view / lax.bitcast_convert_type to int8/int16 keeping bit pattern","Keep Pallas matmul kernels on float32/bfloat16/int8 inputs only"],"exampleFix":"# before\nout = pl.dot(x_u8, w_u8)  # x_u8: uint8\n\n# after\nout = pl.dot(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32))","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\n\ndef ok_for_pallas_dot(*arrays):\n    return all(not jnp.issubdtype(a.dtype, jnp.unsignedinteger) for a in arrays)","typeGuard":"def is_signed_or_float(a) -> bool:\n    import jax.numpy as jnp\n    return not jnp.issubdtype(a.dtype, jnp.unsignedinteger)","tryCatchPattern":null,"preventionTips":["Normalize uint8 inputs to float32/int32 at data-loading time","Add a startup assert on kernel operand dtypes before launching pallas kernels"],"tags":["jax","pallas","triton","gpu","matmul","dtype"],"backgroundTag":"unsupported-dtype-for-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}