{"record":{"id":"fb166782622a96e6","repo":"xai-org/x-algorithm","slug":"q-k-and-v-should-all-be-4d-got-q-ndim-k-n","errorCode":null,"errorMessage":"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}","messagePattern":"q, k, and v should all be 4D, got: (.+?), (.+?), (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":90,"sourceCode":"\n    @property\n    def has_backward_blocks(self) -> bool:\n        return self.block_q_dkv is not None\n\n\ndef _attention_forward(\n    q,\n    k,\n    v,\n    config: TuningConfig,\n    save_residuals: bool = False,\n    bound=None,\n    sm_scale: float = 1.0,\n    cap: float = -1.0,\n    cap_method: str = \"tanh\",\n):\n    if q.ndim != 4 or k.ndim != 4 or v.ndim != 4:\n        raise ValueError(f\"q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.ndim=}\")\n    batch_size, q_seq_len, num_q_heads, head_dim = q.shape\n    _, kv_seq_len, num_kv_heads, _ = k.shape\n    kv_shape = (batch_size, kv_seq_len, num_kv_heads, head_dim)\n    if k.shape != kv_shape:\n        raise ValueError(f\"Expected {k.shape=} to be {kv_shape} (inferred from q)\")\n    if v.shape != kv_shape:\n        raise ValueError(f\"Expected {v.shape=} to be {kv_shape} (inferred from q)\")\n    if (dtype := q.dtype) != k.dtype or dtype != v.dtype:\n        raise ValueError(\n            f\"q, k, and v should all have the same dtype, got: {q.dtype}, {k.dtype}, {v.dtype}\"\n        )\n    if num_q_heads % num_kv_heads:\n        raise ValueError(f\"{num_q_heads=} must be divisible by and {num_kv_heads=}\")\n    q_heads_per_kv_head = num_q_heads // num_kv_heads\n    if head_dim % 64:\n        raise ValueError(f\"{head_dim=} must be divisible by 64\")\n    if jnp.dtype(dtype) not in map(jnp.dtype, [jnp.float16, jnp.bfloat16]):\n        raise NotImplementedError(f\"Only f16 and bf16 are supported, got dtype: {dtype}\")","sourceCodeStart":72,"sourceCodeEnd":108,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L72-L108","documentation":"_attention_forward requires q, k, v in BHSD layout (batch, seq, heads, head_dim) — exactly 4 dimensions each. This is the entry check before any shape inference happens.","triggerScenarios":"Passing 3D inputs (B, S, D without a head axis) or inputs still in (B, H, S, D) when the function expects (B, S, H, D), or accidentally batched extra leading dims.","commonSituations":"Layout confusion between BHSD and BSHD conventions across jax attention implementations; feeding outputs of einops.rearrange with the wrong pattern; passing single-example unbatched tensors.","solutions":["Reshape to 4D (batch, seq_len, num_heads, head_dim)","If you have (B, H, S, D), transpose axes 1 and 2","Add a head dimension of 1 for single-head inputs"],"exampleFix":"# before\nout = attention(q, k, v)  # q is (B, H, S, D)\n# after\nq, k, v = (x.transpose(0, 2, 1, 3) for x in (q, k, v))\nout = attention(q, k, v)","handlingStrategy":"type-guard","validationCode":"null","typeGuard":"def is_bshd4(x) -> bool:\n    return hasattr(x, \"ndim\") and x.ndim == 4","tryCatchPattern":null,"preventionTips":["Standardize on (B, S, H, D) layout project-wide","Write a single pre_attention_check(q, k, v) helper and call it before every attention invocation"],"tags":["jax","pallas","attention","shape-validation"],"backgroundTag":"tensor-shape-mismatch","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}