jax-ml/jax · error · NotImplementedError

float32 top_k is not supported on TPUv3 or older

Error message

float32 top_k is not supported on TPUv3 or older

What it means

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.

Source

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

@register_lowering_rule(lax.top_k_p, ensure_mlir_values=False)
def _top_k_lowering_rule(
    ctx: LoweringRuleContext,
    x,
    *,
    k: int,
    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 = []

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. 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
2. Run on TPUv4 or newer hardware
3. Do the top_k in regular JAX outside the Pallas kernel and pass results in

Example fix

// before
vals, idx = lax.top_k(x_f32, k)  # on TPUv3
// after
vals, idx = lax.top_k(x, k)  # computed outside kernel; kernel only consumes results

from jax._src import tpu_info
def fp32_topk_ok():
    return tpu_info.get_tpu_info().generation >= 4

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• bfloat16 top_k is not supported on TPUv5 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/b99bde93c27767dd. Report an issue: GitHub.