{"record":{"id":"e7517b5690b50b34","repo":"xai-org/x-algorithm","slug":"attn-logit-cap-method-method-r-is-not-supported","errorCode":null,"errorMessage":"attn_logit_cap_method {method!r} is not supported by JaxAttention.","messagePattern":"attn_logit_cap_method (.+?) is not supported by JaxAttention\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"phoenix/xrex/models/attention.py","lineNumber":235,"sourceCode":"    query_input: jax.Array,\n    key_input: jax.Array,\n    pairwise_fn: Callable[..., Any] = jnp.multiply,\n    dtype: Any = jnp.bfloat16,\n):\n    mask = pairwise_fn(jnp.expand_dims(query_input, axis=-1), jnp.expand_dims(key_input, axis=-2))\n    mask = jnp.expand_dims(mask, axis=-3)\n    return mask.astype(dtype)\n\n\ndef _cap_attention_logits(logits: jax.Array, cap: float, method: str) -> jax.Array:\n    if not cap or cap <= 0.0 or method == \"none\":\n        return logits\n    if method == \"tanh\":\n        cap_arr = jnp.array(cap, dtype=logits.dtype)\n        return cap_arr * jnp.tanh(logits / cap_arr)\n    if method == \"soft_sign\":\n        return logits / (1.0 + jnp.abs(logits) / cap)\n    raise ValueError(f\"attn_logit_cap_method {method!r} is not supported by JaxAttention.\")\n\n\nclass JaxAttention(Attention):\n    def call_attn(\n        self,\n        query: jax.Array,\n        key: Optional[jax.Array],\n        value: Optional[jax.Array],\n        segment_ids: Optional[jax.Array],\n        segment_ids_k: Optional[jax.Array],\n        temp: Optional[jax.Array],\n        **kwargs,\n    ):\n        mask = kwargs.get(\"masks\", None)\n        b, t, h, d = query.shape\n        _, _, kv_h, _ = key.shape\n        assert h % kv_h == 0, f\"query_heads {h} must be a multiple of kv_heads {kv_h}\"\n        assert self.sharding_context is not None","sourceCodeStart":217,"sourceCodeEnd":253,"githubUrl":"https://github.com/xai-org/x-algorithm/blob/24c60942c5c5fdad3a6addffb4c6e6d2f228f04f/phoenix/xrex/models/attention.py#L217-L253","documentation":"JaxAttention supports only two attention logit capping methods in _cap_attention_logits: 'tanh' (cap * tanh(logits / cap)) and 'soft_sign' (logits / (1 + |logits| / cap)). Any other value of attn_logit_cap_method reaches the trailing raise and is rejected with the offending value.","triggerScenarios":"Setting config.attn_logit_cap_method to something like 'relu', 'sigmoid', 'none', or a typo like 'tanh ' while using the jax_attn implementation; call_attn invokes _cap_attention_logits when a cap is configured.","commonSituations":"Copying a config from a codebase (e.g. Gemma-style) that accepts more cap methods; typos; version drift where a method was removed or renamed.","solutions":["Use 'tanh' or 'soft_sign' as attn_logit_cap_method.","To disable capping, clear the cap setting rather than passing a sentinel method name.","If you need a new method, add an explicit branch in _cap_attention_logits before the raise."],"exampleFix":"# before\nconfig.attn_logit_cap_method = \"relu\"\n\n# after\nconfig.attn_logit_cap_method = \"tanh\"","handlingStrategy":"validation","validationCode":"assert config.attn_logit_cap_method in {\"tanh\", \"soft_sign\"}, config.attn_logit_cap_method","typeGuard":"def is_valid_cap_method(m: str) -> bool:\n    return m in {\"tanh\", \"soft_sign\"}","tryCatchPattern":null,"preventionTips":["Centralize supported enum values as module-level constants and validate configs against them."],"tags":["attention","logit-cap","enum-validation","config"],"backgroundTag":"invalid-config-value","analyzedSha":"24c60942c5c5fdad3a6addffb4c6e6d2f228f04f","analyzedAt":"2026-08-28T11:40:14.686Z","schemaVersion":2},"datasetVersion":"2026-08-28T16:17:29.566Z"}