{"record":{"id":"8a441f5ba6e9555d","repo":"xai-org/x-algorithm","slug":"keys-must-be-1d-or-2d","errorCode":null,"errorMessage":"keys must be 1D or 2D.","messagePattern":"keys must be 1D or 2D\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/cuda/top_k_by_key/__init__.py","lineNumber":48,"sourceCode":"else:\n    jax.ffi.register_ffi_target(\n        \"xrex_top_k_by_key_radix_select\",\n        fn=top_k_by_key_radix_select_api.top_k_by_key_radix_select(),\n        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),","sourceCodeStart":30,"sourceCodeEnd":66,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/cuda/top_k_by_key/__init__.py#L30-L66","documentation":"top_k_by_key only accepts keys with ndim <= 2 (a single token vector or a batch of vectors); 3D+ arrays (e.g. [batch, seq, vocab]) are rejected because the kernel operates on 1D/2D contiguous key blocks. Flatten or reshape higher-rank score tensors before calling.","triggerScenarios":"Calling top_k_by_key / local_top_k with keys.ndim >= 3, e.g. per-timestep retrieval scores of shape (batch, seq_len, num_items) passed straight from a scoring model.","commonSituations":"Switching from jax.lax.top_k (which allows any rank and reduces over the last axis) to the fused kernel; forgetting to .reshape(-1, n) batch-and-sequence dimensions; new model variants that add a sequence dimension to candidate scoring.","solutions":["Reshape to 2D first: keys2d = keys.reshape(-1, keys.shape[-1]); top_k_by_key(keys2d, k, ...); then reshape results back.","Or loop/vmap over the leading dimension so each call sees at most 2D keys.","For arbitrary-rank needs, keep jax.lax.top_k as a fallback path when keys.ndim > 2."],"exampleFix":"# before\nidx, vals = top_k_by_key(scores, k=k, heuristic_pivot_ratio=0.5)  # scores.shape=(B, S, N)\n\n# after\nB, S, N = scores.shape\nidx, vals = top_k_by_key(scores.reshape(B * S, N).astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)\nidx, vals = idx.reshape(B, S, k), vals.reshape(B, S, k)","handlingStrategy":"type-guard","validationCode":"if keys.ndim > 2:\n    keys = keys.reshape(-1, keys.shape[-1])\nidx, vals = top_k_by_key(keys.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)\n# reshape idx/vals back to (*orig_shape[:-1], k) if needed","typeGuard":"def is_flat_keys(keys: jax.Array) -> bool:\n    return keys.ndim <= 2","tryCatchPattern":"try:\n    idx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)\nexcept ValueError as e:\n    if \"1D or 2D\" in str(e):\n        lead = keys.shape[:-1]\n        idx, vals = top_k_by_key(keys.reshape(-1, keys.shape[-1]).astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)\n        idx, vals = idx.reshape(*lead, -1), vals.reshape(*lead, -1)\n    else:\n        raise","preventionTips":["Centralize top-k calls behind a helper that reshapes and casts once.","For arbitrary-rank tensors, keep jax.lax.top_k as the generic fallback."],"tags":["cuda-kernel","rank-validation","shape","top-k","retrieval"],"backgroundTag":"tensor-rank-not-supported","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}