jax-ml/jax · error · NotImplementedError

Unsigned integer dtype

Error message

Unsigned integer dtype {aval.dtype} is not supported for conv on the Pallas Mosaic TPU backend.

What it means

Dtype guard in the Pallas Mosaic TPU conv lowering: TPU conv hardware ops do not accept unsigned integer operands, so an unsigned dtype on the lhs/rhs raises NotImplementedError; cast to a signed or float type. The input at fault is the conv operand dtype.

Solutions

1. Cast operands to signed dtypes (jnp.int32) before the conv
2. Store quantized data as int8 rather than uint8
3. Use float/bf16 operands for the conv

Example fix

// before
y = lax.conv_general_dilated(x_u8, w_u8, ...)
// after
y = lax.conv_general_dilated(x_u8.astype(jnp.int32), w_u8.astype(jnp.int32), ...)

import jax.numpy as jnp
for a in (lhs, rhs):
    assert not jnp.issubdtype(a.dtype, jnp.unsignedinteger)

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

When it happens

Trigger: lax.conv_general_dilated with uint8/uint16 operands inside a Pallas Mosaic TPU kernel.

Common situations: Quantized inference or image preprocessing pipelines feeding uint8 tensors into conv kernels.

Related errors

• Reductions over unsigned integers not implemented.
• Unsigned integer dtype
• Acc ref dtype must be float32 or int32, got
• {aval_out.dtype}
• {ctx.avals_out[0].dtype}

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

Appendix: source

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

    window_reversal=None,
    feature_group_count=1,
    batch_group_count=1,
    precision=None,
    **_,
):
  if not ctx.is_libtpu_at_least("0.1.0"):
    # When removing this, also remove the pyrefly ignore annotation for ConvOp
    # below.
    raise NotImplementedError("Requires libtpu >= 0.1.0")

  if feature_group_count != 1 or batch_group_count != 1:
    raise NotImplementedError(
        "Grouped convolutions are not supported on Pallas Mosaic TPU backend"
        " yet."
    )
  for aval in ctx.avals_in[:2]:
    if jnp.issubdtype(aval.dtype, jnp.unsignedinteger):
      raise NotImplementedError(
          f"Unsigned integer dtype {aval.dtype} is not supported for conv on"
          " the Pallas Mosaic TPU backend."
      )
  lhs, rhs = args[0], args[1]
  acc = args[2] if len(args) > 2 else None
  (aval_out,) = ctx.avals_out
  out_type = ctx.aval_to_ir_type(aval_out)
  if acc is None:
    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)