xai-org/x-algorithm · error · ValueError

keys must be 1D or 2D.

Error message

keys must be 1D or 2D.

What it means

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.

Source

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

else:
    jax.ffi.register_ffi_target(
        "xrex_top_k_by_key_radix_select",
        fn=top_k_by_key_radix_select_api.top_k_by_key_radix_select(),
        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),

View on GitHub (pinned to 24c60942c5)

Solutions

1. Reshape to 2D first: keys2d = keys.reshape(-1, keys.shape[-1]); top_k_by_key(keys2d, k, ...); then reshape results back.
2. Or loop/vmap over the leading dimension so each call sees at most 2D keys.
3. For arbitrary-rank needs, keep jax.lax.top_k as a fallback path when keys.ndim > 2.

Example fix

# before
idx, vals = top_k_by_key(scores, k=k, heuristic_pivot_ratio=0.5)  # scores.shape=(B, S, N)

# after
B, S, N = scores.shape
idx, vals = top_k_by_key(scores.reshape(B * S, N).astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)
idx, vals = idx.reshape(B, S, k), vals.reshape(B, S, k)

if keys.ndim > 2:
    keys = keys.reshape(-1, keys.shape[-1])
idx, vals = top_k_by_key(keys.astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)
# reshape idx/vals back to (*orig_shape[:-1], k) if needed

def is_flat_keys(keys: jax.Array) -> bool:
    return keys.ndim <= 2

try:
    idx, vals = top_k_by_key(keys, k=k, heuristic_pivot_ratio=0.5)
except ValueError as e:
    if "1D or 2D" in str(e):
        lead = keys.shape[:-1]
        idx, vals = top_k_by_key(keys.reshape(-1, keys.shape[-1]).astype(jnp.bfloat16), k=k, heuristic_pivot_ratio=0.5)
        idx, vals = idx.reshape(*lead, -1), vals.reshape(*lead, -1)
    else:
        raise

When it happens

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

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

Related errors

• Only bfloat16 is supported for keys.
• k ({k}) must be <= n ({n})
• Expected 2D emb_table, got shape={embeddings.shape}
• Unable to create named shape with unnamed dimensions (shape:
• Number of names must match number of dimensions (shape: {sha

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