{"record":{"id":"1752f4bed250da5c","repo":"xai-org/x-algorithm","slug":"expected-bound-to-have-shape-4-got-bound-arr","errorCode":null,"errorMessage":"Expected bound to have shape (4,), got {bound_arr.shape}","messagePattern":"Expected bound to have shape \\(4,\\), got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention_fa3.py","lineNumber":411,"sourceCode":"        grid_names=(\"heads\", \"q_seq\", \"batch\"),\n        num_threads=3,\n        thread_name=\"wg\",\n        compiler_params=plgpu.CompilerParams(approx_math=True),\n    )(q, k, v, bound)\n\n    if save_residuals:\n        assert lse is not None\n        return out, (lse,)\n\n    return out\n\n\ndef _normalize_bound(bound, batch_size: int, seq_len: int):\n    if bound is None:\n        bound = (0, seq_len + 1, seq_len + 1, seq_len + 1)\n    bound_arr = jnp.asarray(bound, dtype=jnp.int32)\n    if bound_arr.shape != (4,):\n        raise ValueError(f\"Expected bound to have shape (4,), got {bound_arr.shape}\")\n    return jnp.broadcast_to(bound_arr, (batch_size, 4))\n\n\n@functools.partial(jax.custom_vjp, nondiff_argnums=(3, 4, 5, 6, 7, 8, 9))\n@functools.partial(\n    jax.jit,\n    static_argnames=[\n        \"config\",\n        \"save_residuals\",\n        \"bound\",\n        \"sm_scale\",\n        \"cap\",\n        \"cap_method\",\n        \"z_loss_weight\",\n    ],\n)\ndef attention(\n    q,","sourceCodeStart":393,"sourceCodeEnd":429,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention_fa3.py#L393-L429","documentation":"_normalize_bound converts the attention bound (segment boundary spec) to an int32 array and requires exactly shape (4,) — the four bounds (start, end, etc.) per batch — before broadcasting to (batch_size, 4). None defaults to (0, seq_len+1, seq_len+1, seq_len+1).","triggerScenarios":"Passing bound as a scalar, a (batch_size, 4) array (double-broadcast attempt), a 3-element list, or nested lists of wrong shape.","commonSituations":"Users pre-broadcasting bounds themselves; passing per-query bounds or Python ints; shape confusion about whether bound is per-batch or global.","solutions":["Pass bound as a length-4 sequence, e.g. [0, end, end, end]","Or pass None to use the default full-range bound","Keep elements ints (they are cast to int32)"],"exampleFix":"# before\nattn = sharded_mha(q, k, v, bound=jnp.array([[0, 128, 128, 128]] * batch))\n# after\nattn = sharded_mha(q, k, v, bound=(0, 128, 128, 128))","handlingStrategy":"type-guard","validationCode":"null","typeGuard":"def is_valid_bound(b) -> bool:\n    return b is None or (hasattr(b, \"__len__\") and len(b) == 4)","tryCatchPattern":null,"preventionTips":["Always pass bound as a plain 4-tuple or None","Let the function do the broadcasting to (batch, 4)"],"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"}