{"record":{"id":"9f0bbe5063896208","repo":"jax-ml/jax","slug":"bfloat16-top-k-is-not-supported-on-tpuv5-or-older","errorCode":null,"errorMessage":"bfloat16 top_k is not supported on TPUv5 or older","messagePattern":"bfloat16 top_k is not supported on TPUv5 or older","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic/lowering.py","lineNumber":3642,"sourceCode":"    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 = []\n    idxs = []\n    curr = operand\n    for _ in range(k):\n      idx = lax.argmax(curr, axis=axis, index_dtype=index_dtype)\n      val = jnp.max(curr, axis=axis)\n      vals.append(val)\n      idxs.append(idx)\n      mask = iota == jnp.expand_dims(idx, axis)","sourceCodeStart":3624,"sourceCodeEnd":3660,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic/lowering.py#L3624-L3660","documentation":"The Mosaic top_k lowering rejects bfloat16 inputs when the target TPU generation is less than 6 (TPUv5 and older). bfloat16 top_k requires TPUv6 (e.g. Trillium-class) hardware.","triggerScenarios":"lax.top_k on bfloat16 arrays inside a Pallas kernel compiled for TPU v4/v5 (tpu generation < 6).","commonSituations":"Porting sampling kernels to v5p/v5e fleets; assuming bfloat16 is universally supported because it is for other ops; CI targeting an older generation.","solutions":["Cast to float32 if hardware is gen >= 4 (v4/v5)","Move top_k outside the kernel on v5 and older","Target TPUv6+ hardware for in-kernel bf16 top_k"],"exampleFix":"// before\nvals, idx = lax.top_k(x_bf16, k)  # on TPUv5\n// after\nvals, idx = lax.top_k(x_bf16.astype(jnp.float32), k)","handlingStrategy":"validation","validationCode":"from jax._src import tpu_info\ndef bf16_topk_ok():\n    return tpu_info.get_tpu_info().generation >= 6","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use float32 top_k on v4/v5; bfloat16 only on v6+","Abstract dtype+generation checks into kernel preconditions"],"tags":["jax","pallas","tpu","top-k","hardware-version","bfloat16"],"backgroundTag":"tpu-generation-unsupported","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}