xai-org/x-algorithm · error · ValueError

cap_method must be in [tanh, soft_sign]

Error message

cap_method must be in [tanh, soft_sign]

What it means

Raised in the dq backward inner loop of the varlen ranker attention kernel when cap > 0.0 and cap_method is invalid. The forward capping options are the only ones with implemented derivatives, so the backward path guards on the same set of values.

Source

Thrown at phoenix/xrex/pallas/ranker_attention_varlen.py:795

            segment_ref,
            (pl.dslice(offset_k, block_k),),
            mask=k_mask,
            other=PADDING_SEGMENT_ID,
        ).astype(jnp.int8)[None, :]
        qk = pl.dot(q, k, trans_b=True)

        if sm_scale != 1.0:
            qk *= sm_scale

        if cap > 0.0:
            if cap_method == "tanh":
                qk_tanh = tanh(qk / cap)
                qk = cap * qk_tanh
            elif cap_method == "soft_sign":
                soft_sign = 1.0 / (1.0 + jnp.abs(qk) / cap)
                qk = qk * soft_sign
            else:
                raise ValueError("cap_method must be in [tanh, soft_sign]")

        mask = jnp.logical_or(seg_k > 0, span_q[:, None] == span_k[None, :])
        mask = jnp.logical_and(seq_q_is_not_padding, mask)
        mask = jnp.logical_and(mask, k_mask[None, :])
        if causal:
            causal_mask = span_q[:, None] >= span_k[None, :]
            mask = jnp.logical_and(mask, causal_mask)

        p = jnp.exp(qk - m[:, None])
        p = jnp.where(mask, p, 0.0)
        dp = pl.dot(do, v, trans_b=True).astype(jnp.float32) - di[:, None]
        ds = p * dp

        if z_loss_weight > 0:
            ds += z_loss_weight * p * ((jnp.log(l + 1e-12) + m) / l)[:, None]

        if cap > 0.0:
            if cap_method == "tanh":

View on GitHub (pinned to 24c60942c5)

Solutions

1. Set cap_method to 'tanh' or 'soft_sign'.
2. Set cap=0.0 to skip capping entirely.
3. Fail fast by validating cap_method at startup in training scripts.

Example fix

# before
grad = jax.grad(loss)(..., cap=50.0, cap_method='softmax_cap')
# after
grad = jax.grad(loss)(..., cap=50.0, cap_method='soft_sign')

assert cap <= 0.0 or cap_method in ('tanh', 'soft_sign')

def valid_cap_for_grad(cap: float, cap_method: str) -> bool:
    return cap <= 0.0 or cap_method in ('tanh', 'soft_sign')

When it happens

Trigger: Running gradients through varlen ranker attention with cap > 0.0 and cap_method other than 'tanh'/'soft_sign'.

Common situations: Bad cap_method only exercised during training; config drift between train and eval paths; sweep jobs hitting the gradient path for the first time.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• Causal attention not supported in the backwards pass yet.
• Need to specify backward blocks.
• {q_seq_len=} must be a multiple of {config.block_q_dq=} * {c
• Invalid backward pass implementation: {backward_pass_impl}
• Invalid backward pass implementation: {backward_pass_impl}

AI-assisted analysis of xai-org/x-algorithm@24c60942c5 (2026-08-28). Data as JSON: /api/errors/cdf8fcfc2bd05e1c. Report an issue: GitHub.