Function softclamp

src/sharp/utils/math.py:75–103

Clamp tensor to min/max in differentiable way. Args: tensor: The tensor to clamp. min: Pair of threshold to start clamping and value to clamp to. The first value should be larger than the second. max: Pair of threshold to start clamping and value to clamp to.

(
    tensor: torch.Tensor,
    min: SoftClampRange | None = None,
    max: SoftClampRange | None = None,
)

def softclamp(
    tensor: torch.Tensor,
    min: SoftClampRange | None = None,
    max: SoftClampRange | None = None,
) -> torch.Tensor:
    """Clamp tensor to min/max in differentiable way.

    Args:
        tensor: The tensor to clamp.
        min: Pair of threshold to start clamping and value to clamp to.
            The first value should be larger than the second.
        max: Pair of threshold to start clamping and value to clamp to.
            The first value should be smaller than the second.

    Returns:
        The clamped tensor.
    """

    def normalize(clamp_range: SoftClampRange) -> torch.Tensor:
        value0, value1 = clamp_range
        return value0 + (value1 - value0) * torch.tanh((tensor - value0) / (value1 - value0))

    tensor_clamped = tensor
    if min is not None:
        tensor_clamped = torch.maximum(tensor_clamped, normalize(min))
    if max is not None:
        tensor_clamped = torch.minimum(tensor_clamped, normalize(max))

    return tensor_clamped


class ClampWithPushback(autograd.Function):