homa 0.2.0__tar.gz → 0.2.93__tar.gz

{homa-0.2.0 → homa-0.2.93}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.0
+Version: 0.2.93
 Summary: A curated list of machine learning and deep learning helpers.
 Author-email: Taha Shieenavaz <tahashieenavaz@gmail.com>
 Requires-Python: >=3.7

{homa-0.2.0 → homa-0.2.93}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "homa"
-version = "0.2.0"
+version = "0.2.93"
 description = "A curated list of machine learning and deep learning helpers."
 authors = [
   { name="Taha Shieenavaz", email="tahashieenavaz@gmail.com" },

homa-0.2.93/src/homa/activations/APLU.py

@@ -0,0 +1,49 @@
+import torch
+
+
+class APLU(torch.nn.Module):
+    def __init__(
+        self, channels: int | None = None, n: int = 2, init_b: str = "linspace"
+    ):
+        super().__init__()
+        self.n = n
+        self.init_b = init_b
+        if channels is None:
+            self.register_parameter("a", None)
+            self.register_parameter("b", None)
+        else:
+            self._init_params(channels, device=None, dtype=None)
+
+    def _init_params(self, channels, device, dtype):
+        a = torch.zeros(channels, self.n, device=device, dtype=dtype)
+        if self.init_b == "linspace":
+            b = (
+                torch.linspace(-1.0, 1.0, steps=self.n, device=device, dtype=dtype)
+                .expand(channels, -1)
+                .contiguous()
+            )
+        else:
+            b = torch.empty(channels, self.n, device=device, dtype=dtype).uniform_(
+                -1.0, 1.0
+            )
+        self.a = torch.nn.Parameter(a)
+        self.b = torch.nn.Parameter(b)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.a is None or self.b is None:
+            self._init_params(x.shape[1], device=x.device, dtype=x.dtype)
+
+        y = F.relu(x)
+        x_exp = x.unsqueeze(-1)
+        expand_shape = (
+            (
+                1,
+                x.shape[1],
+            )
+            + (1,) * (x.dim() - 2)
+            + (self.n,)
+        )
+        a = self.a.view(*expand_shape)
+        b = self.b.view(*expand_shape)
+        hinges = (-x_exp + b).clamp_max(0.0)
+        return y + (a * hinges).sum(dim=-1)

homa-0.2.93/src/homa/activations/ActivationFunction.py

@@ -0,0 +1,6 @@
+import torch
+
+
+class ActivationFunction(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

homa-0.2.93/src/homa/activations/AdaptiveActivationFunction.py

@@ -0,0 +1,15 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class AdaptiveActivationFunction(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def __repr__(self):
+        arguments_text = ""
+
+        if hasattr(self, "num_channels"):
+            arguments_text = f"channels={self.num_channels}"
+
+        return f"{__class__.__name__}({arguments_text})"

homa-0.2.93/src/homa/activations/BaseDLReLU.py

@@ -0,0 +1,34 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class BaseDLReLU(ActivationFunction):
+    def __init__(self, a: float = 0.01, init_mse: float = 1.0, mode: str = "linear"):
+        super().__init__()
+        assert 0.0 < a < 1.0, "a must be in (0,1)"
+        assert mode in ("linear", "exp")
+        self.a = float(a)
+        self.mode = mode
+        self.register_buffer("prev_mse", torch.tensor(float(init_mse)))
+
+    @torch.no_grad()
+    def set_prev_mse(self, mse_value):
+        if isinstance(mse_value, torch.Tensor):
+            mse_value = float(mse_value.detach().cpu().item())
+        self.prev_mse.fill_(mse_value)
+
+    @torch.no_grad()
+    def update_from_loss(self, loss_tensor: torch.Tensor):
+        self.set_prev_mse(loss_tensor)
+
+    def forward(self, z: torch.Tensor, mse_prev: torch.Tensor | float | None = None):
+        b_t = (
+            self.prev_mse
+            if mse_prev is None
+            else (torch.as_tensor(mse_prev, device=z.device, dtype=z.dtype))
+        )
+        if self.mode == "linear":
+            slope = self.a * b_t
+        else:
+            slope = self.a * torch.exp(-b_t)
+        return torch.where(z >= 0, z, slope * z)

homa-0.2.93/src/homa/activations/CaLU.py

@@ -0,0 +1,13 @@
+import torch
+import math
+from .ActivationFunction import ActivationFunction
+
+
+class CaLU(ActivationFunction):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a = torch.arctan(x) / math.pi
+        b = 0.5
+        return x * (a + b)

homa-0.2.93/src/homa/activations/DLReLU.py

@@ -0,0 +1,6 @@
+from .BaseDLReLU import BaseDLReLU
+
+
+class DLReLU(BaseDLReLU):
+    def __init__(self, a: float = 0.01, init_mse: float = 1.0):
+        super().__init__(a=a, init_mse=init_mse, mode="linear")

homa-0.2.93/src/homa/activations/ERF.py

@@ -0,0 +1,10 @@
+import torch
+
+
+class ERF(torch.nn.Module):
+    def __init__(self, alpha=1.0):
+        super().__init__()
+        self.alpha = torch.nn.Parameter(torch.tensor(alpha))
+
+    def forward(self, x: torch.Tensor):
+        return x * torch.erf(self.alpha * x)

homa-0.2.93/src/homa/activations/Elliot.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class Elliot(ActivationFunction):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return 0.5 + torch.div(0.5 * x, 1 + torch.abs(x))

homa-0.2.93/src/homa/activations/ExpExpish.py

@@ -0,0 +1,9 @@
+import torch
+
+
+class ExpExpish(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.exp(-torch.exp(-x))

homa-0.2.93/src/homa/activations/ExponentialDLReLU.py

@@ -0,0 +1,6 @@
+from .BaseDLReLU import BaseDLReLU
+
+
+class ExponentialDLReLU(BaseDLReLU):
+    def __init__(self, a: float = 0.01, init_mse: float = 1.0):
+        super().__init__(a=a, init_mse=init_mse, mode="exp")

homa-0.2.93/src/homa/activations/ExponentialSwish.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class ExponentialSwish(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.exp(-x) * torch.sigmoid(x)

homa-0.2.93/src/homa/activations/GCU.py

@@ -0,0 +1,9 @@
+import torch
+
+
+class GCU(torch.nn.Module):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor):
+        return x * torch.cos(x)

homa-0.2.93/src/homa/activations/GaLU.py

@@ -0,0 +1,11 @@
+from .GaussianReLU import GaussianReLU
+
+
+class GaLU(GaussianReLU):
+    def __init__(
+        self,
+        channels: int | None = None,
+        max_input: float = 1.0,
+    ):
+        self.hats = [(2.0, 2.0), (1.0, 1.0), (3.0, 1.0)]
+        super().__init__(self.hats, channels=channels, max_input=max_input)

homa-0.2.93/src/homa/activations/GaussianReLU.py

@@ -0,0 +1,50 @@
+import torch
+from typing import Sequence, Tuple
+
+
+class GaussianReLU(torch.nn.Module):
+    def __init__(
+        self,
+        alphas_lambdas: Sequence[Tuple[float, float]],
+        channels: int | None = None,
+        max_input: float = 1.0,
+    ):
+        super().__init__()
+        self.M = float(max_input)
+        self.register_buffer(
+            "alphas", torch.tensor([a for a, _ in alphas_lambdas], dtype=torch.float32)
+        )
+        self.register_buffer(
+            "lambdas", torch.tensor([l for _, l in alphas_lambdas], dtype=torch.float32)
+        )
+        self.K = len(alphas_lambdas)
+
+        if channels is None:
+            self.register_parameter("c0", None)  # per-channel (PReLU slope)
+            self.register_parameter("c", None)  # (C, K) coefficients
+        else:
+            self._init_params(channels, None, None)
+
+    def _init_params(self, C: int, device, dtype):
+        self.c0 = torch.nn.Parameter(torch.zeros(C, device=device, dtype=dtype))
+        self.c = torch.nn.Parameter(torch.zeros(C, self.K, device=device, dtype=dtype))
+
+    def _expand_param(p: torch.Tensor, x: torch.Tensor, add_K: bool = False):
+        shape = (
+            (1, x.shape[1]) + (1,) * (x.dim() - 2) + ((p.shape[-1],) if add_K else ())
+        )
+        return p.view(shape)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.c0 is None or self.c is None:
+            self._init_params(x.shape[1], x.device, x.dtype)
+        c0 = self._expand_param(self.c0, x)
+        y = torch.nn.functional.relu(x) - c0 * torch.nn.functional.relu(-x)
+        a = self.alphas.to(x.device, x.dtype).view(*((1,) * x.dim()), -1)
+        l = self.lambdas.to(x.device, x.dtype).view(*((1,) * x.dim()), -1)
+        xE = x.unsqueeze(-1)
+        term1 = (l * self.M - (xE - a * self.M).abs()).clamp_min(0.0)
+        term2 = ((xE - a * self.M - 2 * l * self.M).abs() - l * self.M).clamp_max(0.0)
+        hats = term1 + term2
+        c = self._expand_param(self.c, x, add_K=True)  # (1,C,...,K)
+        return y + (c * hats).sum(dim=-1)

homa-0.2.93/src/homa/activations/GeneralizedSwish.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class GeneralizedSwish(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x * torch.sigmoid(torch.exp(-x))

homa-0.2.93/src/homa/activations/Gish.py

@@ -0,0 +1,11 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class Gish(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        a = -torch.exp(x)
+        return x * torch.log(2 - torch.exp(a))

homa-0.2.93/src/homa/activations/LaLU.py

@@ -0,0 +1,11 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class LaLU(ActivationFunction):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        phi_laplace = torch.where(x >= 0, 1 - 0.5 * torch.exp(-x), 0.5 * torch.exp(x))
+        return x * phi_laplace

homa-0.2.93/src/homa/activations/LogLogish.py

@@ -0,0 +1,10 @@
+import torch
+
+
+class LogLogish(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        a = -torch.exp(x)
+        return x * (1 - torch.exp(a))

homa-0.2.93/src/homa/activations/LogSigmoid.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class LogSigmoid(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        return torch.log(torch.sigmoid(x))

homa-0.2.93/src/homa/activations/Logish.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class Logish(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.log1p(torch.sigmoid(x))

homa-0.2.93/src/homa/activations/MeLU.py

@@ -0,0 +1,11 @@
+from .MexicanReLU import MexicanReLU
+
+
+class MeLU(MexicanReLU):
+    def __init__(self, channels: int | None = None, max_input: float = 1.0):
+        self.hats = [
+            (2.0, 2.0),
+            (1.0, 1.0),
+            (3.0, 1.0),
+        ]
+        super().__init__(self.hats, channels=channels, max_input=max_input)

homa-0.2.93/src/homa/activations/MexicanReLU.py

@@ -0,0 +1,49 @@
+import torch
+from typing import Sequence, Tuple
+
+
+class MexicanReLU(torch.nn.Module):
+    def __init__(
+        self,
+        alphas_lambdas: Sequence[Tuple[float, float]],
+        channels: int | None = None,
+        max_input: float = 1.0,
+    ):
+        super().__init__()
+        self.M = float(max_input)
+        self.register_buffer(
+            "alphas", torch.tensor([a for a, _ in alphas_lambdas], dtype=torch.float32)
+        )
+        self.register_buffer(
+            "lambdas", torch.tensor([l for _, l in alphas_lambdas], dtype=torch.float32)
+        )
+        self.K = len(alphas_lambdas)
+
+        if channels is None:
+            self.register_parameter("c0", None)  # PReLU negative slope (per-channel)
+            self.register_parameter("c", None)  # (C, K) coefficients
+        else:
+            self._init_params(channels, device=None, dtype=None)
+
+    def _init_params(self, C: int, device, dtype):
+        self.c0 = torch.nn.Parameter(torch.zeros(C, device=device, dtype=dtype))
+        self.c = torch.nn.Parameter(torch.zeros(C, self.K, device=device, dtype=dtype))
+
+    def _expand_param(p: torch.Tensor, x: torch.Tensor, n_extra: int = 0):
+        shape = (
+            (1, x.shape[1]) + (1,) * (x.dim() - 2) + ((p.shape[-1],) if n_extra else ())
+        )
+        return p.view(shape)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.c0 is None or self.c is None:
+            self._init_params(x.shape[1], x.device, x.dtype)
+        c0 = self._expand_param(self.c0, x)
+        y = F.relu(x) - c0 * F.relu(-x)
+        xE = x.unsqueeze(-1)
+        cE = self._expand_param(self.c, x, n_extra=1)
+        aE = self.alphas.to(x.device, x.dtype).view(*((1,) * x.dim()), -1)  # (..., K)
+        lE = self.lambdas.to(x.device, x.dtype).view(*((1,) * x.dim()), -1)  # (..., K)
+        hats = (lE * self.M - (xE - aE * self.M).abs()).clamp_min(0.0)
+        y = y + (cE * hats).sum(dim=-1)
+        return y

homa-0.2.93/src/homa/activations/MinSin.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class MinSin(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.min(x, torch.sin(x))

homa-0.2.93/src/homa/activations/NReLU.py

@@ -0,0 +1,12 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class NReLU(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        sigma = x.std(unbiased=False)
+        a = torch.randn_like(x) * sigma
+        return torch.where(x >= 0, x + a, 0)

homa-0.2.93/src/homa/activations/NoisyReLU.py

@@ -0,0 +1,6 @@
+from .NReLU import NReLU
+
+
+class NoisyReLU(NReLU):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

homa-0.2.93/src/homa/activations/PLogish.py

@@ -0,0 +1,6 @@
+from .ParametricLogish import ParametricLogish
+
+
+class PLogish(ParametricLogish):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)

homa-0.2.93/src/homa/activations/ParametricLogish.py

@@ -0,0 +1,13 @@
+import torch
+
+
+class ParametricLogish(torch.nn.Module):
+    def __init__(self, alpha: float = 1.0, beta: float = 10.0):
+        super().__init__()
+        self.alpha = alpha
+        self.beta = beta
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a = torch.sigmoid(self.beta * x)
+        b = torch.log(1 + a)
+        return self.alpha * x * b

homa-0.2.93/src/homa/activations/Phish.py

@@ -0,0 +1,11 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class Phish(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a = torch.nn.functional.gelu(x)
+        return x * torch.tanh(a)

homa-0.2.93/src/homa/activations/RReLU.py

@@ -0,0 +1,16 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class RReLU(ActivationFunction):
+    def __init__(self, lower: int = 3, upper: int = 8):
+        super(RReLU, self).__init__()
+        self.lower = lower
+        self.upper = upper
+
+    def forward(self, x):
+        if self.training:
+            a = torch.empty_like(x).uniform_(self.lower, self.upper)
+        else:
+            a = (self.lower + self.upper) / 2.0
+        return torch.where(x >= 0, x, x / a)

homa-0.2.93/src/homa/activations/RandomizedSlopedReLU.py

@@ -0,0 +1,7 @@
+import random
+from .SlopedReLU import SlopedReLU
+
+
+class RandomizedSlopedReLU(SlopedReLU):
+    def __init__(self):
+        super().__init__(alpha=random.uniform(1, 10))

homa-0.2.93/src/homa/activations/SGELU.py

@@ -0,0 +1,12 @@
+import torch
+import math
+from .ActivationFunction import ActivationFunction
+
+
+class SGELU(ActivationFunction):
+    def __init__(self, alpha: float = 0.1, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.alpha = alpha
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.alpha * x * torch.erf(x / math.sqrt(2))

homa-0.2.93/src/homa/activations/SReLU.py

@@ -0,0 +1,37 @@
+import torch
+
+
+class SReLU(torch.nn.Module):
+    def __init__(self, channels: int | None = None, max_input: float = 1.0):
+        super().__init__()
+        self.M = float(max_input)
+        if channels is None:
+            self.register_parameter("t_l", None)
+            self.register_parameter("t_r", None)
+            self.register_parameter("a_l", None)
+            self.register_parameter("a_r", None)
+        else:
+            self._init_params(channels, None, None)
+
+    def _init_params(self, C: int, device, dtype):
+        self.t_l = torch.nn.Parameter(torch.zeros(C, device=device, dtype=dtype))
+        self.t_r = torch.nn.Parameter(
+            torch.full((C,), self.M, device=device, dtype=dtype)
+        )
+        self.a_l = torch.nn.Parameter(torch.zeros(C, device=device, dtype=dtype))
+        self.a_r = torch.nn.Parameter(torch.ones(C, device=device, dtype=dtype))
+
+    def _expand_param(p: torch.Tensor, x: torch.Tensor):
+        return p.view((1, x.shape[1]) + (1,) * (x.dim() - 2))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.t_l is None:
+            self._init_params(x.shape[1], x.device, x.dtype)
+
+        t_l = self._expand_param(self.t_l, x)
+        t_r = self._expand_param(self.t_r, x)
+        a_l = self._expand_param(self.a_l, x)
+        a_r = self._expand_param(self.a_r, x)
+        y = torch.where(x < t_l, t_l + a_l * (x - t_l), x)
+        y = torch.where(x > t_r, t_r + a_r * (x - t_r), y)
+        return y

homa-0.2.93/src/homa/activations/SelfArctan.py

@@ -0,0 +1,9 @@
+import torch
+
+
+class SelfArctan(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return x * torch.arctan(x)

homa-0.2.93/src/homa/activations/ShiftedReLU.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class ShiftedReLU(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        return torch.max(-1, x)

homa-0.2.93/src/homa/activations/SigmoidDerivative.py

@@ -0,0 +1,10 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class SigmoidDerivative(ActivationFunction):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.exp(-x) * torch.sigmoid(x).pow(2)

homa-0.2.93/src/homa/activations/SineReLU.py

@@ -0,0 +1,11 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class SineReLU(ActivationFunction):
+    def __init__(self, epsilon: float = 0.0025):
+        super().__init__()
+        self.epsilon: float = epsilon
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.where(x >= 0, x, self.epsilon * (torch.sin(x) - torch.cos(x)))

homa-0.2.93/src/homa/activations/SlopedReLU.py

@@ -0,0 +1,13 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class SlopedReLU(ActivationFunction):
+    def __init__(self, alpha: float):
+        super().__init__()
+        if not (1 <= alpha < 10):
+            raise Exception("Sloped ReLU slope needs to be in [1, 10].")
+        self.alpha: float = alpha
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return torch.where(x >= 0, self.alpha * x, 0)

homa-0.2.93/src/homa/activations/SmallGaLU.py

@@ -0,0 +1,11 @@
+from .GaussianReLU import GaussianReLU
+
+
+class SmallGaLU(GaussianReLU):
+    def __init__(
+        self,
+        channels: int | None = None,
+        max_input: float = 1.0,
+    ):
+        self.hats = [(2.0, 2.0)]
+        super().__init__(self.hats, channels=channels, max_input=max_input)

homa-0.2.93/src/homa/activations/Smish.py

@@ -0,0 +1,9 @@
+import torch
+
+
+class Smish(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor):
+        return x * torch.tanh(torch.log1p(torch.sigmoid(x)))

homa-0.2.93/src/homa/activations/SoftsignRReLU.py

@@ -0,0 +1,17 @@
+import torch
+from .ActivationFunction import ActivationFunction
+
+
+class SoftsignRReLU(ActivationFunction):
+    def __init__(self, lower: float = 1 / 8, upper: float = 1 / 3, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lower = lower
+        self.upper = upper
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        a = torch.empty_like(x).uniform_(self.lower, self.upper)
+        return torch.where(
+            x >= 0,
+            torch.div(1, (1 + x).pow(2)) + x,
+            torch.div(1, (1 + x).pow(2)) + a * x,
+        )