{"record":{"id":"392834cf3b187b38","repo":"jax-ml/jax","slug":"reductions-over-unsigned-integers-not-implemented","errorCode":null,"errorMessage":"Reductions over unsigned integers not implemented.","messagePattern":"Reductions over unsigned integers not implemented\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2620,"sourceCode":"        val = val[jnp.newaxis, ...]\n        axes = [axis + 1 for axis in axes]\n        val = reduce_fn(val, axis=axes, keepdims=True)\n        # Squeeze lowers to vector.ExtractOp which will place the final\n        # value in a scalar register.\n        return jnp.squeeze(val)\n      proxy_lowering = lower_fun(_proxy_fun)\n      return proxy_lowering(ctx, x, axes=axes)\n\n    if jnp.issubdtype(x_aval.dtype, jnp.floating):\n      kind = type_to_kind[jnp.floating]\n      val = type_to_identity[jnp.floating]\n      val = ir.FloatAttr.get(ctx.aval_to_ir_type(x_aval, shape=()), val)\n    elif x_aval.dtype == jnp.int32:\n      kind = type_to_kind[jnp.signedinteger]\n      val = type_to_identity[jnp.signedinteger]\n      val = ir.IntegerAttr.get(ir.IntegerType.get_signless(32), val)\n    elif jnp.issubdtype(x_aval.dtype, jnp.unsignedinteger):\n      raise NotImplementedError(\n          \"Reductions over unsigned integers not implemented.\"\n      )\n    else:\n      raise NotImplementedError(\n          f\"Reductions over {x_aval.dtype} not implemented.\")\n    out_type = ctx.aval_to_ir_type(ctx.avals_out[0])\n    identity = ir.DenseElementsAttr.get_splat(out_type, val)\n    acc = arith.constant(out_type, identity)\n    return vector.multi_reduction(kind, x, acc, axes)\n  return _lowering_rule\n\n\nREDUCE_MAX_KINDS = {\n    jnp.floating: vector.CombiningKind.MAXIMUMF,\n    jnp.signedinteger: vector.CombiningKind.MAXSI,\n    jnp.unsignedinteger: vector.CombiningKind.MAXUI,\n}\nREDUCE_MAX_IDENTITY = {","sourceCodeStart":2602,"sourceCodeEnd":2638,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2602-L2638","documentation":"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.","triggerScenarios":"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.","commonSituations":"Reducing indices, hashes, or bitmask data stored as uint32; loading uint8 data and reducing without cast.","solutions":["Cast to a supported dtype before reducing: x.astype(jnp.int32) or jnp.float32","For max/min of unsigned values where sign matters, use int32 with careful range handling or f32 if values fit"],"exampleFix":"# before\ntotal = jnp.sum(x_uint32, axis=0)\n# after\ntotal = jnp.sum(x_uint32.astype(jnp.int32), axis=0)","handlingStrategy":"type-guard","validationCode":"if jnp.issubdtype(x.dtype, jnp.unsignedinteger):\n    x = x.astype(jnp.int32)\ntotal = jnp.sum(x, axis=0)","typeGuard":"def reduction_dtype_ok(dtype) -> bool:\n    import jax.numpy as jnp\n    return (jnp.issubdtype(dtype, jnp.floating)\n            or dtype == jnp.int32\n            or jnp.issubdtype(dtype, jnp.complexfloating))","tryCatchPattern":null,"preventionTips":["Cast to int32/float32 before reductions in Pallas kernels","Keep a kernel-side dtype policy: compute in f32/i32 only"],"tags":["jax","pallas","tpu","reduction","unsigned-integer","dtype"],"backgroundTag":"unsigned-integer-reduction-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}