Logit Soft-Capping

Gemma 2 and Grok soft-cap logits so no value runs away: $\text{softcap}(z) = c \cdot \tanh(z / c)$. Small logits pass through almost unchanged, while huge ones saturate smoothly at ±c (unlike hard clipping, the gradient never becomes exactly zero).

Implement:

• softcap(logits, cap)
• softcap_grad(logits, cap): the derivative $1 - \tanh^2(z / c)$
• max_prob(logits): the largest softmax probability (computed stably)
• capped_ratio(logits, cap): max_prob(softcap(logits, cap)) / max_prob(logits), showing how capping tempers over-confident predictions