jax-ml/jax · error · NotImplementedError

bfloat16 top_k is not supported on TPUv5 or older

Error message

bfloat16 top_k is not supported on TPUv5 or older

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic/lowering.py:3642

    axis: int,
    is_stable: bool,
):
  input_dtype = ctx.avals_in[0].dtype
  if input_dtype not in (jnp.float32, jnp.bfloat16):
    raise NotImplementedError(
        f"Pallas top_k only supports float32 and bfloat16, got {input_dtype}"
    )
  tpu_gen = tpu_info.get_tpu_info().generation
  if input_dtype == jnp.float32 and tpu_gen < 4:
    raise NotImplementedError(
        "float32 top_k is not supported on TPUv3 or older"
    )
  if input_dtype == jnp.bfloat16 and tpu_gen < 6:
    raise NotImplementedError(
        "bfloat16 top_k is not supported on TPUv5 or older"
    )
  if is_stable:
    raise NotImplementedError(
        "is_stable=True is not supported in Pallas top_k. For efficiency, only"
        " is_stable=False is supported"
    )

  def _top_k_impl(operand, *, k: int, axis: int = -1):
    axis = axis % operand.ndim
    index_dtype = jnp.int16 if operand.dtype == jnp.bfloat16 else jnp.int32
    iota = lax.broadcasted_iota(index_dtype, operand.shape, axis)
    min_val = jnp.finfo(operand.dtype).min
    vals = []
    idxs = []
    curr = operand
    for _ in range(k):
      idx = lax.argmax(curr, axis=axis, index_dtype=index_dtype)
      val = jnp.max(curr, axis=axis)
      vals.append(val)
      idxs.append(idx)
      mask = iota == jnp.expand_dims(idx, axis)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast to float32 if hardware is gen >= 4 (v4/v5)
2. Move top_k outside the kernel on v5 and older
3. Target TPUv6+ hardware for in-kernel bf16 top_k

Example fix

// before
vals, idx = lax.top_k(x_bf16, k)  # on TPUv5
// after
vals, idx = lax.top_k(x_bf16.astype(jnp.float32), k)

from jax._src import tpu_info
def bf16_topk_ok():
    return tpu_info.get_tpu_info().generation >= 6

When it happens

Trigger: lax.top_k on bfloat16 arrays inside a Pallas kernel compiled for TPU v4/v5 (tpu generation < 6).

Common situations: Porting sampling kernels to v5p/v5e fleets; assuming bfloat16 is universally supported because it is for other ops; CI targeting an older generation.

Related errors

• float32 top_k is not supported on TPUv3 or older
• Pallas top_k only supports float32 and bfloat16, got {input_
• TPU version must be 4 or higher.
• Compiler params for platform {platform} cannot be used for {
• Memory space {self.memory_space} is not supported by mesh {s

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/9f0bbe5063896208. Report an issue: GitHub.