xai-org/x-algorithm · error · ValueError

cap_method must be in [tanh, soft_sign], got {cap_method}

Error message

cap_method must be in [tanh, soft_sign], got {cap_method}

What it means

The ranker attention forward kernel caps logits (qk) before the softmax using either tanh or soft_sign squashing. If cap_method is not one of these two strings, the kernel body raises this ValueError inside the Pallas program.

Source

Thrown at phoenix/xrex/pallas/ranker_attention.py:94

        seg_k = pl.load(
            segment_ref,
            (pl.dslice(start_k * block_k, block_k),),
        )
        temp = pl.load(temp_ref, (pl.dslice(start_q * block_q, block_q),))
        temp = jnp.expand_dims(temp, axis=-1)
        mask = jnp.equal(jnp.zeros_like(seg_q), jnp.expand_dims(seg_k, axis=-2))
        qk = jnp.zeros([block_q, block_k], dtype=jnp.float32)
        qk += pl.dot(q, k)
        if sm_scale != 1.0:
            qk *= sm_scale

        if cap > 0.0:
            if cap_method == "tanh":
                qk = cap * tanh(qk / cap)
            elif cap_method == "soft_sign":
                qk = qk / (1.0 + jnp.abs(qk) / cap)
            else:
                raise ValueError(f"cap_method must be in [tanh, soft_sign], got {cap_method}")
        qk = qk * temp
        span_q = start_q * block_q + jnp.arange(block_q)
        span_k = start_k * block_k + jnp.arange(block_k)
        if causal:
            causal_mask = span_q[:, None] >= (span_k[None, :] + inverted_sliding_window_sizep1)
            mask = jnp.logical_and(causal_mask, mask)
        mask = jnp.logical_or(mask, span_q[:, None] == span_k[None, :])
        if window_len > 0:
            window_mask = span_k[None, :] > span_q[:, None] - window_len
            mask = jnp.logical_and(mask, window_mask)

        padding_mask = jnp.logical_and(seg_q != 2, jnp.expand_dims(seg_k != 2, axis=-2))
        mask = jnp.logical_and(mask, padding_mask)

        qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
        m_curr = jnp.maximum(jnp.max(qk, axis=1), m_prev)
        l_new = jnp.exp(m_prev - m_curr)
        l_new = jax.lax.select(jnp.isnan(l_new), jnp.ones_like(l_new), l_new)

View on GitHub (pinned to 24c60942c5)

Solutions

1. Set cap_method to 'tanh' or 'soft_sign'
2. If you want no capping, set cap=0.0 so the capping branch is skipped entirely
3. Add config validation for cap_method before launching the kernel (Pallas errors inside kernels are hard to debug)

Example fix

# before
attn_out = ranker_attention(q, k, v, cap=10.0, cap_method="clamp")

# after
attn_out = ranker_attention(q, k, v, cap=10.0, cap_method="tanh")

if cap > 0.0:
    assert cap_method in ("tanh", "soft_sign"), f"bad cap_method={cap_method!r}"

def valid_cap_method(m: str) -> bool:
    return m in ("tanh", "soft_sign")

When it happens

Trigger: Calling the ranker attention kernel with cap > 0.0 and cap_method set to an unrecognized value (e.g. 'sigmoid', 'clip', or a typo). Note cap_method == 0 disables capping only when cap <= 0.0.

Common situations: Config typos, copying config from another attention implementation with different capping names, or defaulting cap_method to None while setting cap > 0.

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

• cap_method must be in [tanh, soft_sign]
• Invalid backward pass implementation: {backward_pass_impl}
• Invalid backward pass implementation: {backward_pass_impl}
• cap_method must be in [tanh, soft_sign, none], got {cap_meth
• q, k, and v should all be 4D, got: {q.ndim=}, {k.ndim=}, {v.

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