import torch


def soft_thresholding(x, b, a=None):
    """Remap values between [-a, b] to 0, keep the rest linear.
    """
    if a is None:
        a = b
    return (torch.clamp(x - b, min=0) * (x > 0) +
            torch.clamp(x + a, max=0) * (x <= 0))