jax-ml/jax · error · NotImplementedError

Unsigned integer dtype {aval.dtype} is not supported for con

Error message

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

What it means

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).

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)

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 {aval.dtype} is not supported for dot
• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over {x_aval.dtype} not implemented.

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