{"record":{"id":"0881fc97966a2d28","repo":"xai-org/x-algorithm","slug":"cap-method-must-be-in-tanh-soft-sign","errorCode":null,"errorMessage":"cap_method must be in [tanh, soft_sign]","messagePattern":"cap_method must be in \\[tanh, soft_sign\\]","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/pallas/ranker_attention.py","lineNumber":451,"sourceCode":"            segment_ref,\n            (pl.dslice(start_k * block_k, block_k),),\n        )\n        mask = jnp.equal(jnp.zeros_like(seg_q), jnp.expand_dims(seg_k, axis=-2))\n        temp = pl.load(temp_ref, (pl.dslice(start_q * block_q, block_q),))\n        temp = jnp.expand_dims(temp, axis=-1)\n        qk = jnp.zeros((block_q, block_k), dtype=jnp.float32)\n        qk += pl.dot(q, k.T)\n        if sm_scale != 1.0:\n            qk *= sm_scale\n        if cap > 0.0:\n            if cap_method == \"tanh\":\n                qk_tanh = tanh(qk / cap)\n                qk = cap * qk_tanh\n            elif cap_method == \"soft_sign\":\n                soft_sign = 1.0 / (1.0 + jnp.abs(qk) / cap)\n                qk = qk * soft_sign\n            else:\n                raise ValueError(\"cap_method must be in [tanh, soft_sign]\")\n        qk *= temp\n        span_k = start_k * block_k + jnp.arange(block_k)\n        if causal:\n            causal_mask = span_q[:, None] >= span_k[None, :] + inverted_sliding_window_sizep1\n            mask = jnp.logical_and(causal_mask, mask)\n        mask = jnp.logical_or(mask, span_q[:, None] == span_k[None, :])\n        if window_len > 0:\n            window_mask = span_k[None, :] > span_q[:, None] - window_len\n            mask = jnp.logical_and(mask, window_mask)\n        qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)\n        p = jnp.exp(qk - m[:, None])\n        dp = jnp.zeros((block_q, block_k), dtype=jnp.float32) - di[:, None]\n        dp = dp + pl.dot(do, v.T)\n        ds = p * dp\n        if z_loss_weight > 0:\n            ds += z_loss_weight * p * ((jnp.log(l + 1e-12) + m) / l)[:, None]\n        ds *= temp\n        if cap > 0.0:","sourceCodeStart":433,"sourceCodeEnd":469,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/pallas/ranker_attention.py#L433-L469","documentation":"Backward kernel twin of the forward capping check: while recomputing the forward pass inside the backward inner_loop, the kernel hits an unknown cap_method. The gradient of the cap function differs for tanh (1 - tanh^2) vs soft_sign, so the method must be known.","triggerScenarios":"Running the vjp/grad of ranker attention with cap > 0.0 and a cap_method string other than 'tanh' or 'soft_sign' — typically the same bad value that would fail in forward.","commonSituations":"A cap_method typo that only manifests during training (when backward runs), or a forward-only test that passed while the training run fails.","solutions":["Fix cap_method to 'tanh' or 'soft_sign' at the call site / config","Set cap=0.0 to disable capping if unsure","Validate cap_method in Python before entering the jitted/vjp path so the failure is raised outside the kernel"],"exampleFix":"# before\nloss = grad(loss_fn)(params)  # loss_fn uses cap_method=\"sigmoid\"\n\n# after\nloss = grad(loss_fn)(params)  # loss_fn uses cap_method=\"soft_sign\"","handlingStrategy":"validation","validationCode":"assert cap == 0.0 or cap_method in (\"tanh\", \"soft_sign\")","typeGuard":"def valid_cap_config(cap: float, cap_method: str) -> bool:\n    return cap <= 0.0 or cap_method in (\"tanh\", \"soft_sign\")","tryCatchPattern":null,"preventionTips":["Validate kwargs once in Python before jit; errors inside Pallas kernels are opaque","Run grad checks in CI for custom attention configs"],"tags":["attention","pallas","backward","logit-capping","invalid-argument"],"backgroundTag":"invalid-enum-value","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}