def soft_cap(values: torch.Tensor, cap_value: float | torch.Tensor | None) -> torch.Tensor:
    """
    Soft caps a tensor to a value.

    Performs a tanh operation on the logits and scales the result to the cap value. Common technique in attention
    and output language heads to prevent large logits from dominating the softmax. See for example Gemma2:
    https://arxiv.org/abs/2408.00118

    Args:
        values: The tensor to cap.
        cap_value: The value to cap the values to. If None, no cap is applied.

    Returns:
        The capped values.
    """
    if cap_value is None:
        return values
    return cap_value * torch.tanh(values / cap_value)

    Soft caps a tensor to a value.

    Performs a tanh operation on the logits and scales the result to the cap value. Common technique in attention
    and output language heads to prevent large logits from dominating the softmax. See for example Gemma2:
    https://arxiv.org/abs/2408.00118

    Args:
        values: The tensor to cap.
        cap_value: The value to cap the values to. If None, no cap is applied.

    Returns:
        The capped values.
    
xlstm/xlstm_large/components.py