# AUTOGENERATED! DO NOT EDIT! File to edit: ../../nbs/utils/07_clampling.ipynb.

# %% auto 0
__all__ = ['symmetric_tanh_clamp', 'floor_softplus_clamp', 'ceil_softplus_clamp']

# %% ../../nbs/utils/07_clampling.ipynb 2
import torch
from torch.nn import functional as F

# %% ../../nbs/utils/07_clampling.ipynb 5
def symmetric_tanh_clamp(x: torch.tensor, max_value: float) -> torch.tensor:
    """
    Clamps the input tensor to the range [-max_value, max_value] using tanh while keeping it differentiable.

    Parameters:
        x: The input tensor to be clamped.
        max_value: The maximum absolute value to clamp to.

    Returns:
        Clamped differentiable tensor.
    """
    # Ensure max value is positive and greater than 0
    max_value = abs(max_value)
    if max_value == 0:
        raise ValueError("max_value must be greater than 0")

    # Clamp
    clamped_x = max_value * torch.tanh(x / max_value)
    return clamped_x

# %% ../../nbs/utils/07_clampling.ipynb 7
def floor_softplus_clamp(x: torch.tensor, min_value: float) -> torch.tensor:
    """
    Clamps the input tensor to be greater than or equal to min_value using softplus while keeping it differentiable.

    Parameters:
        x: The input tensor to be clamped.
        min_value: The minimum value to clamp to.

    Returns:
        Clamped differentiable tensor.
    """
    # Clamp
    clamped_x = min_value + F.softplus(x - min_value)
    return clamped_x

# %% ../../nbs/utils/07_clampling.ipynb 9
def ceil_softplus_clamp(x: torch.tensor, max_value: float) -> torch.tensor:
    """
    Clamps the input tensor to be less than or equal to max_value using softplus while keeping it differentiable.

    Parameters:
        x: The input tensor to be clamped.
        max_value: The maximum value to clamp to.

    Returns:
        Clamped differentiable tensor.
    """
    # Clamp
    clamped_x = max_value - F.softplus(max_value - x)
    return clamped_x
