{"record":{"id":"3586cf118effa725","repo":"jax-ml/jax","slug":"unsigned-integer-dtype-aval-dtype-is-not-support-3586cf","errorCode":null,"errorMessage":"Unsigned integer dtype {aval.dtype} is not supported for conv on the Pallas Mosaic TPU backend.","messagePattern":"Unsigned integer dtype (.+?) is not supported for conv on the Pallas Mosaic TPU backend\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3058,"sourceCode":"    window_reversal=None,\n    feature_group_count=1,\n    batch_group_count=1,\n    precision=None,\n    **_,\n):\n  if not ctx.is_libtpu_at_least(\"0.1.0\"):\n    # When removing this, also remove the pyrefly ignore annotation for ConvOp\n    # below.\n    raise NotImplementedError(\"Requires libtpu >= 0.1.0\")\n\n  if feature_group_count != 1 or batch_group_count != 1:\n    raise NotImplementedError(\n        \"Grouped convolutions are not supported on Pallas Mosaic TPU backend\"\n        \" yet.\"\n    )\n  for aval in ctx.avals_in[:2]:\n    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):\n      raise NotImplementedError(\n          f\"Unsigned integer dtype {aval.dtype} is not supported for conv on\"\n          \" the Pallas Mosaic TPU backend.\"\n      )\n  lhs, rhs = args[0], args[1]\n  acc = args[2] if len(args) > 2 else None\n  (aval_out,) = ctx.avals_out\n  out_type = ctx.aval_to_ir_type(aval_out)\n  if acc is None:\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":3040,"sourceCodeEnd":3076,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3040-L3076","documentation":"Mosaic's conv lowering rejects unsigned integer operand dtypes (like dot, TPU contract ops interpret integers as signed), checking the first two input avals (lhs/rhs).","triggerScenarios":"lax.conv_general_dilated with uint8/uint16 operands inside a Pallas Mosaic TPU kernel.","commonSituations":"Quantized inference or image preprocessing pipelines feeding uint8 tensors into conv kernels.","solutions":["Cast operands to signed dtypes (jnp.int32) before the conv","Store quantized data as int8 rather than uint8","Use float/bf16 operands for the conv"],"exampleFix":"// before\ny = lax.conv_general_dilated(x_u8, w_u8, ...)\n// after\ny = lax.conv_general_dilated(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32), ...)","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nfor a in (lhs, rhs):\n    assert not jnp.issubdtype(a.dtype, jnp.unsignedinteger)","typeGuard":"def conv_operand_ok(dt) -> bool:\n    import jax.numpy as jnp\n    return not jnp.issubdtype(dt, jnp.unsignedinteger)","tryCatchPattern":null,"preventionTips":["Cast quantized uint data to int32 before entering kernels","Prefer signed or float dtypes end-to-end on TPU"],"tags":["jax","pallas","tpu","convolution","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"}