vendored model

sharp/utils/math.py

@@ -0,0 +1,183 @@
+"""Contains utility math functions.
+
+For licensing see accompanying LICENSE file.
+Copyright (C) 2025 Apple Inc. All Rights Reserved.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Callable, Literal, NamedTuple, Tuple, Union
+
+import torch
+from torch import autograd
+
+ActivationType = Literal[
+    "linear",
+    "exp",
+    "sigmoid",
+    "softplus",
+    "relu_with_pushback",
+    "hard_sigmoid_with_pushback",
+]
+ActivationFunction = Callable[[torch.Tensor], torch.Tensor]
+
+
+class ActivationPair(NamedTuple):
+    """A pair of forward and inverse activation functions."""
+
+    forward: ActivationFunction
+    inverse: ActivationFunction
+
+
+def create_activation_pair(activation_type: ActivationType) -> ActivationPair:
+    """Create activation function and corresponding inverse function.
+
+    Args:
+        activation_type: The activation type to create.
+
+    Returns:
+        The corresponding activation functions and the corresponding inverse function.
+    """
+    if activation_type == "linear":
+        return ActivationPair(lambda x: x, lambda x: x)
+    elif activation_type == "exp":
+        return ActivationPair(torch.exp, torch.log)
+    elif activation_type == "sigmoid":
+        return ActivationPair(torch.sigmoid, inverse_sigmoid)
+    elif activation_type == "softplus":
+        return ActivationPair(torch.nn.functional.softplus, inverse_softplus)
+    elif activation_type == "relu_with_pushback":
+        return ActivationPair(relu_with_pushback, lambda x: x)
+    elif activation_type == "hard_sigmoid_with_pushback":
+        return ActivationPair(hard_sigmoid_with_pushback, lambda x: 6.0 * x - 3.0)
+    else:
+        raise ValueError(f"Unsupported activation function: {activation_type}.")
+
+
+def inverse_sigmoid(tensor: torch.Tensor) -> torch.Tensor:
+    """Compute inverse sigmoid."""
+    return torch.log(tensor / (1.0 - tensor))
+
+
+def inverse_softplus(tensor: torch.Tensor, eps: float = 1e-06) -> torch.Tensor:
+    """Compute inverse softplus."""
+    tensor = tensor.clamp_min(eps)
+    sigmoid = torch.sigmoid(-tensor)
+    exp = sigmoid / (1.0 - sigmoid)
+    return tensor + torch.log(-exp + 1.0)
+
+
+# The first value describes the threshold from where clamping will be applied, while
+# the second value describes the value to clamp with.
+SoftClampRange = Tuple[Union[torch.Tensor, float], Union[torch.Tensor, float]]
+
+
+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):
+    """Implementation of clamp_with_pushback function."""
+
+    @staticmethod
+    def forward(
+        ctx: Any,
+        tensor: torch.Tensor,
+        min: float | None,
+        max: float | None,
+        pushback: float,
+    ) -> torch.Tensor:
+        """Apply clamp."""
+        if min is not None and max is not None and min >= max:
+            raise ValueError("Only min < max is supported.")
+
+        ctx.save_for_backward(tensor)
+        ctx.min = min
+        ctx.max = max
+        ctx.pushback = pushback
+        return torch.clamp(tensor, min=min, max=max)
+
+    @staticmethod
+    def backward(  # type: ignore[override] # Deal with buggy torch annotations.
+        ctx: Any, grad_in: torch.Tensor
+    ) -> tuple[torch.Tensor, None, None, None]:
+        """Compute gradient of clamp with pushback."""
+        grad_out = grad_in.clone()
+        (tensor,) = ctx.saved_tensors
+
+        if ctx.min is not None:
+            mask_min = tensor < ctx.min
+            grad_out[mask_min] = -ctx.pushback
+
+        if ctx.max is not None:
+            mask_max = tensor > ctx.max
+            grad_out[mask_max] = ctx.pushback
+
+        return grad_out, None, None, None
+
+
+def clamp_with_pushback(
+    tensor: torch.Tensor,
+    min: float | None = None,
+    max: float | None = None,
+    pushback: float = 1e-2,
+) -> torch.Tensor:
+    """Variant of clamp function which avoid the vanishing gradient problem.
+
+    This function is equivalent to adding a regularizer of the form
+
+        pushback * sum_i (
+            relu(min - preactivation_i) + relu(preactivation_i - max)
+        )
+
+    to the full loss function, which pushes clamped values back.
+
+    When used in minimization problems, pushback should be greater than
+    zero. In maximization problems, pushback should be smaller than zero.
+    """
+    output = ClampWithPushback.apply(tensor, min, max, pushback)
+    assert isinstance(output, torch.Tensor)
+    return output
+
+
+def hard_sigmoid_with_pushback(x: torch.Tensor, slope: float = 1.0 / 6.0) -> torch.Tensor:
+    """Apply hard sigmoid with pushback.
+
+    For compatibility reasons, we follow the default PyTorch implementation with a
+    default slope of 1/6:
+
+        https://pytorch.org/docs/stable/generated/torch.nn.Hardsigmoid.html
+    """
+    return clamp_with_pushback(slope * x + 0.5, min=0.0, max=1.0)
+
+
+def relu_with_pushback(x: torch.Tensor) -> torch.Tensor:
+    """Compute relu with pushback."""
+    return clamp_with_pushback(x, min=0.0)