xai-org/x-algorithm · error · ValueError

k ({k}) must be <= n ({n})

Error message

k ({k}) must be <= n ({n})

What it means

top_k_by_key validates that the requested k does not exceed n, where n = keys.shape[-1] (the last dimension of the keys array). Requesting more top elements than exist along that axis is meaningless, so it fails fast with a ValueError before dispatching to the radix-select, async, or default CUDA API.

Source

Thrown at phoenix/xrex/cuda/top_k_by_key/__init__.py:52

        platform="CUDA",
    )


def top_k_by_key(
    keys: jax.Array,
    k: int,
    heuristic_pivot_ratio: float,
    use_async: bool = False,
    use_radix_select: bool = False,
):
    if keys.dtype != jnp.bfloat16:
        raise ValueError("Only bfloat16 is supported for keys.")
    if keys.ndim > 2:
        raise ValueError("keys must be 1D or 2D.")

    n = keys.shape[-1]
    if k > n:
        raise ValueError(f"k ({k}) must be <= n ({n})")

    if use_radix_select:
        api = top_k_by_key_radix_select_api
    elif use_async:
        api = top_k_by_key_async_api
    else:
        api = top_k_by_key_api
    if api is None or jax.default_backend() != "gpu":
        sorted_keys, sorted_indices = jax.lax.top_k(keys, k)
        return sorted_keys, sorted_indices.astype(jnp.int32)

    out_shape = (k,) if keys.ndim == 1 else (keys.shape[0], k)
    out_types = [
        jax.ShapeDtypeStruct(shape=out_shape, dtype=keys.dtype),
        jax.ShapeDtypeStruct(shape=out_shape, dtype=jnp.int32),
    ]
    if use_radix_select:
        call = jax.ffi.ffi_call(

View on GitHub (pinned to 24c60942c5)

Solutions

1. Reduce k so that k <= keys.shape[-1]
2. If k is derived dynamically, clamp it: k = min(k, keys.shape[-1])
3. Verify you are reading n from the last axis of keys, not from another tensor's shape

Example fix

// before
topk_vals, topk_idx = local_top_k(keys, k=512)  # keys.shape[-1] == 256
// after
k = min(512, keys.shape[-1])
topk_vals, topk_idx = local_top_k(keys, k=k)

n = keys.shape[-1]
assert k <= n, f"k={k} exceeds n={n} along keys.shape[-1]"
# or clamp: k = min(k, n)

When it happens

Trigger: Calling top_k_by_key(keys, k, ...) (directly or via local_top_k) where k > keys.shape[-1], e.g. keys of shape (B, 128) with k=256, or passing a per-batch k scalar sized against the wrong axis.

Common situations: Hardcoded k not adjusted when sequence length/dimension shrinks; computing k from a different tensor than keys (e.g. using values.shape[0] instead of keys.shape[-1]); off-by-one from k = n + 1 in loop sweeps over k.

Related errors

• Only int32 is supported for unique.
• Only 1D arrays are supported for unique.
• valid_block_upper and valid_block_lower must be provided tog
• type checking expression %s failed: invalid argument type: e
• ASTNode %s expected %d arguments, %d passed.

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/f5a0c53ed174db40. Report an issue: GitHub.