{"record":{"id":"f0c3cd00d35f648d","repo":"jax-ml/jax","slug":"reductions-over-x-aval-dtype-not-implemented","errorCode":null,"errorMessage":"Reductions over {x_aval.dtype} not implemented.","messagePattern":"Reductions over (.+?) not implemented\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":2624,"sourceCode":"        # 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 = {\n    jnp.floating: float(\"-inf\"),\n    jnp.signedinteger: np.iinfo(np.int32).min,\n}\n_reduce_max_lowering_rule = reduce_lowering_rule(","sourceCodeStart":2606,"sourceCodeEnd":2642,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L2606-L2642","documentation":"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.","triggerScenarios":"jnp.sum/jnp.max/jnp.min/etc. on int8/int16/int64/bool/unsigned arrays inside a TPU Pallas kernel.","commonSituations":"Summing quantized int8 activations, bool counts, or int64 indices in a kernel without casting to int32/float32 first.","solutions":["Cast operands to jnp.int32 or jnp.float32 before the reduction","For int64, check whether values fit in int32 or restructure the kernel to avoid 64-bit reductions"],"exampleFix":"# before\ns = jnp.sum(mask_bool)\n# after\ns = jnp.sum(mask_bool.astype(jnp.int32))","handlingStrategy":"type-guard","validationCode":"SUPPORTED = (jnp.float32, jnp.float16, jnp.bfloat16, jnp.int32)\nif x.dtype not in SUPPORTED and not jnp.issubdtype(x.dtype, jnp.floating):\n    x = x.astype(jnp.float32)\nout = 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":["Always cast int8/int16/int64/bool to int32 or f32 before reducing in TPU kernels","Add dtype assertions at kernel entry"],"tags":["jax","pallas","tpu","reduction","dtype"],"backgroundTag":"unsupported-dtype-reduction","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}