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

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.

Source

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

            segment_ref,
            (pl.dslice(start_k * block_k, block_k),),
        )
        mask = jnp.equal(jnp.zeros_like(seg_q), jnp.expand_dims(seg_k, axis=-2))
        temp = pl.load(temp_ref, (pl.dslice(start_q * block_q, block_q),))
        temp = jnp.expand_dims(temp, axis=-1)
        qk = jnp.zeros((block_q, block_k), dtype=jnp.float32)
        qk += pl.dot(q, k.T)
        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]")
        qk *= temp
        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)
        qk = jnp.where(mask, qk, DEFAULT_MASK_VALUE)
        p = jnp.exp(qk - m[:, None])
        dp = jnp.zeros((block_q, block_k), dtype=jnp.float32) - di[:, None]
        dp = dp + pl.dot(do, v.T)
        ds = p * dp
        if z_loss_weight > 0:
            ds += z_loss_weight * p * ((jnp.log(l + 1e-12) + m) / l)[:, None]
        ds *= temp
        if cap > 0.0:

View on GitHub (pinned to 24c60942c5)

Solutions

1. Fix cap_method to 'tanh' or 'soft_sign' at the call site / config
2. Set cap=0.0 to disable capping if unsure
3. Validate cap_method in Python before entering the jitted/vjp path so the failure is raised outside the kernel

Example fix

# before
loss = grad(loss_fn)(params)  # loss_fn uses cap_method="sigmoid"

# after
loss = grad(loss_fn)(params)  # loss_fn uses cap_method="soft_sign"

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

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

When it happens

Trigger: 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.

Common situations: A cap_method typo that only manifests during training (when backward runs), or a forward-only test that passed while the training run fails.

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], got {cap_method}
• 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/0881fc97966a2d28. Report an issue: GitHub.