{"record":{"id":"f5a0c53ed174db40","repo":"xai-org/x-algorithm","slug":"k-k-must-be-n-n","errorCode":null,"errorMessage":"k ({k}) must be <= n ({n})","messagePattern":"k \\((.+?)\\) must be <= n \\((.+?)\\)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/top_k_by_key/__init__.py","lineNumber":52,"sourceCode":"        platform=\"CUDA\",\n    )\n\n\ndef top_k_by_key(\n    keys: jax.Array,\n    k: int,\n    heuristic_pivot_ratio: float,\n    use_async: bool = False,\n    use_radix_select: bool = False,\n):\n    if keys.dtype != jnp.bfloat16:\n        raise ValueError(\"Only bfloat16 is supported for keys.\")\n    if keys.ndim > 2:\n        raise ValueError(\"keys must be 1D or 2D.\")\n\n    n = keys.shape[-1]\n    if k > n:\n        raise ValueError(f\"k ({k}) must be <= n ({n})\")\n\n    if use_radix_select:\n        api = top_k_by_key_radix_select_api\n    elif use_async:\n        api = top_k_by_key_async_api\n    else:\n        api = top_k_by_key_api\n    if api is None or jax.default_backend() != \"gpu\":\n        sorted_keys, sorted_indices = jax.lax.top_k(keys, k)\n        return sorted_keys, sorted_indices.astype(jnp.int32)\n\n    out_shape = (k,) if keys.ndim == 1 else (keys.shape[0], k)\n    out_types = [\n        jax.ShapeDtypeStruct(shape=out_shape, dtype=keys.dtype),\n        jax.ShapeDtypeStruct(shape=out_shape, dtype=jnp.int32),\n    ]\n    if use_radix_select:\n        call = jax.ffi.ffi_call(","sourceCodeStart":34,"sourceCodeEnd":70,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/top_k_by_key/__init__.py#L34-L70","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Reduce k so that k <= keys.shape[-1]","If k is derived dynamically, clamp it: k = min(k, keys.shape[-1])","Verify you are reading n from the last axis of keys, not from another tensor's shape"],"exampleFix":"// before\ntopk_vals, topk_idx = local_top_k(keys, k=512)  # keys.shape[-1] == 256\n// after\nk = min(512, keys.shape[-1])\ntopk_vals, topk_idx = local_top_k(keys, k=k)","handlingStrategy":"validation","validationCode":"n = keys.shape[-1]\nassert k <= n, f\"k={k} exceeds n={n} along keys.shape[-1]\"\n# or clamp: k = min(k, n)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Derive k from keys.shape[-1] at the call site instead of hardcoding","Clamp dynamic k values with min(k, keys.shape[-1])","Add shape/k assertions in test fixtures for top_k_by_key"],"tags":["cuda","jax","top-k","argument-validation"],"backgroundTag":"argument-out-of-range","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}