jax-ml/jax · error · NotImplementedError

Pallas top_k only supports float32 and bfloat16, got {input_

Error message

Pallas top_k only supports float32 and bfloat16, got {input_dtype}

What it means

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.

Source

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

def _argmin_lowering_rule(ctx: LoweringRuleContext, x, axes, index_dtype):
  return _reduce_index_helper(
      ctx, x, axes, index_dtype,
      ir.Attribute.parse("#tpu.reduction_kind<arg_min>")
  )


@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):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Cast input to jnp.float32 before top_k (cast indices/values back after)
2. Use bfloat16 if precision loss is acceptable
3. Implement top-k manually via sorting (lax.sort) if the dtype must be preserved
4. Move top_k outside the Pallas kernel

Example fix

// before
vals, idx = lax.top_k(scores, k)  # scores is float16 in kernel
// after
vals, idx = lax.top_k(scores.astype(jnp.float32), k)
vals = vals.astype(scores.dtype)

import jax.numpy as jnp
def topk_dtype_ok(dt):
    return dt in (jnp.float32, jnp.bfloat16)

def topk_supported(dt): return jnp.dtype(dt) in (jnp.dtype(jnp.float32), jnp.dtype(jnp.bfloat16))

When it happens

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

Common situations: top_k over float16 logits or integer scores inside a sampling kernel; converting an existing GPU Pallas top_k kernel to TPU.

Related errors

• Acc ref dtype must be float32 or int32, got {dtype}
• masked swap with non-32-bit data
• Reductions over unsigned integers not implemented.
• Reductions over {x_aval.dtype} not implemented.
• Unsigned integer dtype {aval.dtype} is not supported for dot

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