jax-ml/jax · error · NotImplementedError

Reductions over {x_aval.dtype} not implemented.

Error message

Reductions over {x_aval.dtype} not implemented.

What it means

Raised by the TPU vector reduction lowering for dtypes outside the supported set (floats, int32, complex). Any other dtype (int8, int16, int64, bool, unsigned) reaching vector.multi_reduction lowering raises this NotImplementedError.

Source

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

        # 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 = {
    jnp.floating: float("-inf"),
    jnp.signedinteger: np.iinfo(np.int32).min,
}
_reduce_max_lowering_rule = reduce_lowering_rule(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast operands to jnp.int32 or jnp.float32 before the reduction
2. For int64, check whether values fit in int32 or restructure the kernel to avoid 64-bit reductions

Example fix

# before
s = jnp.sum(mask_bool)
# after
s = jnp.sum(mask_bool.astype(jnp.int32))

SUPPORTED = (jnp.float32, jnp.float16, jnp.bfloat16, jnp.int32)
if x.dtype not in SUPPORTED and not jnp.issubdtype(x.dtype, jnp.floating):
    x = x.astype(jnp.float32)
out = 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: jnp.sum/jnp.max/jnp.min/etc. on int8/int16/int64/bool/unsigned arrays inside a TPU Pallas kernel.

Common situations: Summing quantized int8 activations, bool counts, or int64 indices in a kernel without casting to int32/float32 first.

Related errors

• Reductions over unsigned integers not implemented.
• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Unsigned integer dtype {aval.dtype} is not supported for dot
• {ctx.avals_out[0].dtype}

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