jax-ml/jax · error · NotImplementedError

Reductions over unsigned integers not implemented.

Error message

Reductions over unsigned integers not implemented.

What it means

Raised by the TPU vector reduction lowering when the reduced value has an unsigned integer dtype. The mapping from dtype to reduction kind only implements floats, signed int32, and complex; unsigned integer reductions are not implemented in this path.

Source

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

        val = val[jnp.newaxis, ...]
        axes = [axis + 1 for axis in axes]
        val = reduce_fn(val, axis=axes, keepdims=True)
        # Squeeze lowers to vector.ExtractOp which will place the final
        # value in a scalar register.
        return jnp.squeeze(val)
      proxy_lowering = lower_fun(_proxy_fun)
      return proxy_lowering(ctx, x, axes=axes)

    if jnp.issubdtype(x_aval.dtype, jnp.floating):
      kind = type_to_kind[jnp.floating]
      val = type_to_identity[jnp.floating]
      val = ir.FloatAttr.get(ctx.aval_to_ir_type(x_aval, shape=()), val)
    elif x_aval.dtype == jnp.int32:
      kind = type_to_kind[jnp.signedinteger]
      val = type_to_identity[jnp.signedinteger]
      val = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), val)
    elif jnp.issubdtype(x_aval.dtype, jnp.unsignedinteger):
      raise NotImplementedError(
          "Reductions over unsigned integers not implemented."
      )
    else:
      raise NotImplementedError(
          f"Reductions over {x_aval.dtype} not implemented.")
    out_type = ctx.aval_to_ir_type(ctx.avals_out[0])
    identity = ir.DenseElementsAttr.get_splat(out_type, val)
    acc = arith.constant(out_type, identity)
    return vector.multi_reduction(kind, x, acc, axes)
  return _lowering_rule


REDUCE_MAX_KINDS = {
    jnp.floating: vector.CombiningKind.MAXIMUMF,
    jnp.signedinteger: vector.CombiningKind.MAXSI,
    jnp.unsignedinteger: vector.CombiningKind.MAXUI,
}
REDUCE_MAX_IDENTITY = {

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to a supported dtype before reducing: x.astype(jnp.int32) or jnp.float32
2. For max/min of unsigned values where sign matters, use int32 with careful range handling or f32 if values fit

Example fix

# before
total = jnp.sum(x_uint32, axis=0)
# after
total = jnp.sum(x_uint32.astype(jnp.int32), axis=0)

if jnp.issubdtype(x.dtype, jnp.unsignedinteger):
    x = x.astype(jnp.int32)
total = jnp.sum(x, axis=0)

def reduction_dtype_ok(dtype) -> bool:
    import jax.numpy as jnp
    return (jnp.issubdtype(dtype, jnp.floating)
            or dtype == jnp.int32
            or jnp.issubdtype(dtype, jnp.complexfloating))

When it happens

Trigger: Calling jnp.sum/max/min (or any reduction primitive lowered via vector.multi_reduction) on unsigned arrays (uint8/uint16/uint32/uint64) inside a TPU Pallas kernel.

Common situations: Reducing indices, hashes, or bitmask data stored as uint32; loading uint8 data and reducing without cast.

Related errors

• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot
• Unsigned integer dtype {aval.dtype} is not supported for con
• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data

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