{"record":{"id":"694567df3fa0fdab","repo":"xai-org/x-algorithm","slug":"invalid-backward-pass-implementation-backward-pa-694567","errorCode":null,"errorMessage":"Invalid backward pass implementation: {backward_pass_impl}","messagePattern":"Invalid backward pass implementation: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention.py","lineNumber":781,"sourceCode":"                ),\n            ],\n            out_specs=[\n                pl.BlockSpec(\n                    index_map=lambda j, k, _: (j, 0, k, 0),\n                    block_shape=(None, seq_len, None, head_dim),\n                ),\n                pl.BlockSpec(\n                    index_map=lambda j, k, _: (j, 0, k, 0),\n                    block_shape=(None, seq_len, None, head_dim),\n                ),\n            ],\n            name=\"mha_backward_kv\",\n            debug=debug,\n            interpret=interpret,\n            compiler_params=CompilerParams(num_warps=num_warps, num_stages=2),\n        )(q, k, v, temp, segment_ids, out, do_scaled, l, m, delta)\n    else:\n        raise ValueError(f\"Invalid backward pass implementation: {backward_pass_impl}\")\n    return dq.astype(q.dtype), dk, dv, dtemp, dsegment\n\n\nmha.defvjp(_mha_forward, _mha_backward)\n\n\ndef mha_reference(\n    q,\n    k,\n    v,\n    sm_scale=1.0,\n    cap=-1.0,\n    cap_method=\"tanh\",\n    inverted_sliding_window_sizep1: int = 0,\n):\n    logits = jnp.einsum(\"bqhc,bkhc->bhqk\", q, k).astype(jnp.float32)\n    logits *= sm_scale\n    if cap > 0.0:","sourceCodeStart":763,"sourceCodeEnd":799,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention.py#L763-L799","documentation":"Raised in _mha_backward when the backward_pass_impl argument is not one of the supported implementation names. The code dispatches between pallas kernels (e.g. mha_backward_kv) based on a string selector, and any unrecognized string falls into the else branch.","triggerScenarios":"Calling mha (with defvjp) so JAX triggers the VJP, while passing a backward_pass_impl string other than the supported kernel names (e.g. a typo like 'kv_nosplit' or 'flash').","commonSituations":"Typos in config, or code written against an older/newer version where the set of backward implementations changed; also passing None or empty string explicitly.","solutions":["Check the valid backward_pass_impl values in _mha_backward's if/elif chain (phoenix/xrex/pallas/ranker_attention.py around line 770-781)","Pass one of the supported implementation names, e.g. 'kv'","If you need a different impl, upgrade/downgrade to the version that supports it"],"exampleFix":"// before\nout = mha(q, k, v, backward_pass_impl=\"kv_split\")\n// after\nout = mha(q, k, v, backward_pass_impl=\"kv\")","handlingStrategy":"validation","validationCode":"valid_bwd = {\"kv\"}  # extend from the if/elif chain in _mha_backward\nassert backward_pass_impl in valid_bwd, f\"unsupported {backward_pass_impl=}\"","typeGuard":"null","tryCatchPattern":null,"preventionTips":["Centralize backward_pass_impl strings as constants instead of raw strings","Add a unit test asserting each config's backward impl is supported"],"tags":["jax","pallas","attention","invalid-argument"],"backgroundTag":"invalid-enum-argument","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}