{"record":{"id":"b99bde93c27767dd","repo":"jax-ml/jax","slug":"float32-top-k-is-not-supported-on-tpuv3-or-older","errorCode":null,"errorMessage":"float32 top_k is not supported on TPUv3 or older","messagePattern":"float32 top_k is not supported on TPUv3 or older","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3634,"sourceCode":"\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):\n    axis = axis % operand.ndim\n    index_dtype = jnp.int16 if operand.dtype == jnp.bfloat16 else jnp.int32\n    iota = lax.broadcasted_iota(index_dtype, operand.shape, axis)\n    min_val = jnp.finfo(operand.dtype).min\n    vals = []","sourceCodeStart":3616,"sourceCodeEnd":3652,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3616-L3652","documentation":"The Mosaic top_k lowering rejects float32 inputs when the target TPU generation is less than 4 (TPUv3 and older). float32 top_k hardware/libsupport only exists from TPUv4 onward.","triggerScenarios":"Compiling a Pallas kernel containing lax.top_k on float32 while targeting TPU v2/v3 (tpu_info generation < 4), e.g. on legacy TPUv3 hardware or an emulator configured to v3.","commonSituations":"Running legacy TPUv3 hardware with modern Pallas code; CI emulators defaulting to an older TPU generation; kernels developed on v5 failing when reused on older fleets.","solutions":["Cast top_k input to bfloat16 (supported on gen >= 6 only — check next guard) — for v3, compute top_k outside the kernel on the host","Run on TPUv4 or newer hardware","Do the top_k in regular JAX outside the Pallas kernel and pass results in"],"exampleFix":"// before\nvals, idx = lax.top_k(x_f32, k)  # on TPUv3\n// after\nvals, idx = lax.top_k(x, k)  # computed outside kernel; kernel only consumes results","handlingStrategy":"validation","validationCode":"from jax._src import tpu_info\ndef fp32_topk_ok():\n    return tpu_info.get_tpu_info().generation >= 4","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Check TPU generation before compiling kernels with float32 top_k","Keep a host-side fallback path for top_k on old hardware"],"tags":["jax","pallas","tpu","top-k","hardware-version"],"backgroundTag":"tpu-generation-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}