{"record":{"id":"871ac8e23c57b7c4","repo":"jax-ml/jax","slug":"pallas-top-k-only-supports-float32-and-bfloat16-g","errorCode":null,"errorMessage":"Pallas top_k only supports float32 and bfloat16, got {input_dtype}","messagePattern":"Pallas top_k only supports float32 and bfloat16, got (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3629,"sourceCode":"def _argmin_lowering_rule(ctx: LoweringRuleContext, x, axes, index_dtype):\n  return _reduce_index_helper(\n      ctx, x, axes, index_dtype,\n      ir.Attribute.parse(\"#tpu.reduction_kind<arg_min>\")\n  )\n\n\n@register_lowering_rule(lax.top_k_p, ensure_mlir_values=False)\ndef _top_k_lowering_rule(\n    ctx: LoweringRuleContext,\n    x,\n    *,\n    k: int,\n    axis: int,\n    is_stable: bool,\n):\n  input_dtype = ctx.avals_in[0].dtype\n  if input_dtype not in (jnp.float32, jnp.bfloat16):\n    raise NotImplementedError(\n        f\"Pallas top_k only supports float32 and bfloat16, got {input_dtype}\"\n    )\n  tpu_gen = tpu_info.get_tpu_info().generation\n  if input_dtype == jnp.float32 and tpu_gen < 4:\n    raise NotImplementedError(\n        \"float32 top_k is not supported on TPUv3 or older\"\n    )\n  if input_dtype == jnp.bfloat16 and tpu_gen < 6:\n    raise NotImplementedError(\n        \"bfloat16 top_k is not supported on TPUv5 or older\"\n    )\n  if is_stable:\n    raise NotImplementedError(\n        \"is_stable=True is not supported in Pallas top_k. For efficiency, only\"\n        \" is_stable=False is supported\"\n    )\n\n  def _top_k_impl(operand, *, k: int, axis: int = -1):","sourceCodeStart":3611,"sourceCodeEnd":3647,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3611-L3647","documentation":"Mosaic's top_k lowering only supports float32 and bfloat16 input dtypes. Any other dtype (float16, integers, etc.) raises NotImplementedError with the offending dtype in the message.","triggerScenarios":"Calling jax.lax.top_k (or jnp.top_k equivalents traced into the kernel) on non-float32/bfloat16 arrays inside a jax.pallas TPU kernel.","commonSituations":"top_k over float16 logits or integer scores inside a sampling kernel; converting an existing GPU Pallas top_k kernel to TPU.","solutions":["Cast input to jnp.float32 before top_k (cast indices/values back after)","Use bfloat16 if precision loss is acceptable","Implement top-k manually via sorting (lax.sort) if the dtype must be preserved","Move top_k outside the Pallas kernel"],"exampleFix":"// before\nvals, idx = lax.top_k(scores, k)  # scores is float16 in kernel\n// after\nvals, idx = lax.top_k(scores.astype(jnp.float32), k)\nvals = vals.astype(scores.dtype)","handlingStrategy":"type-guard","validationCode":"import jax.numpy as jnp\ndef topk_dtype_ok(dt):\n    return dt in (jnp.float32, jnp.bfloat16)","typeGuard":"def topk_supported(dt): return jnp.dtype(dt) in (jnp.dtype(jnp.float32), jnp.dtype(jnp.bfloat16))","tryCatchPattern":null,"preventionTips":["Cast scores to float32 before top_k in kernels","Add a dtype assertion helper in kernel test suites"],"tags":["jax","pallas","tpu","top-k","dtype"],"backgroundTag":"unsupported-dtype-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}