homa 0.2.0__tar.gz → 0.2.1__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.0
+Version: 0.2.1
 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.1}/pyproject.toml

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

homa-0.2.1/src/homa/activations/classes/APLU.py

@@ -0,0 +1,53 @@
+import torch
+from torch import nn
+
+
+class APLU(nn.Module):
+    def __init__(self, channels: int, max_input: float = 1.0):
+        super().__init__()
+        if channels <= 0:
+            raise ValueError(f"Number of channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.max_input = float(max_input)
+        self.alpha = nn.Parameter(torch.empty(self.channels))
+        self.beta = nn.Parameter(torch.empty(self.channels))
+        self.gamma = nn.Parameter(torch.empty(self.channels))
+        self.xi = nn.Parameter(torch.empty(self.channels))
+        self.psi = nn.Parameter(torch.empty(self.channels))
+        self.mu = nn.Parameter(torch.empty(self.channels))
+        self.reset_parameters()
+
+    def reset_parameters(self) -> None:
+        with torch.no_grad():
+            self.alpha.zero_()
+            self.beta.zero_()
+            self.gamma.zero_()
+            self.xi.uniform_(0.0, self.max_input)
+            self.psi.uniform_(0.0, self.max_input)
+            self.mu.uniform_(0.0, self.max_input)
+
+    @staticmethod
+    def _reshape_for_broadcast(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        return param.view(1, param.shape[0], *([1] * (ref.ndim - 2)))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.ndim < 2:
+            raise ValueError(
+                f"APLU expects inputs with at least two dimensions (N, C, ...), got {tuple(x.shape)}."
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"APLU was initialized with {self.channels} channels but received input with {int(x.shape[1])}."
+            )
+
+        alpha = self._reshape_for_broadcast(self.alpha, x)
+        beta = self._reshape_for_broadcast(self.beta, x)
+        gamma = self._reshape_for_broadcast(self.gamma, x)
+        xi = self._reshape_for_broadcast(self.xi, x)
+        psi = self._reshape_for_broadcast(self.psi, x)
+        mu = self._reshape_for_broadcast(self.mu, x)
+        output = torch.relu(x)
+        output = output + alpha * torch.relu(-x + xi)
+        output = output + beta * torch.relu(-x + psi)
+        output = output + gamma * torch.relu(-x + mu)
+        return output

homa-0.2.1/src/homa/activations/classes/GALU.py

@@ -0,0 +1,61 @@
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F
+
+
+class GALU(nn.Module):
+    def __init__(self, channels: int, max_input: float = 1.0):
+        super().__init__()
+        if max_input <= 0:
+            raise ValueError("max_input must be positive.")
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.max_input = float(max_input)
+        self.alpha = Parameter(torch.empty(self.channels))
+        self.beta = Parameter(torch.empty(self.channels))
+        self.gamma = Parameter(torch.empty(self.channels))
+        self.delta = Parameter(torch.empty(self.channels))
+        self.reset_parameters()
+
+    @staticmethod
+    def _reshape(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        return param.view(1, param.shape[0], *([1] * (ref.ndim - 2)))
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            for param in (self.alpha, self.beta, self.gamma, self.delta):
+                param.zero_()
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim < 2:
+            raise ValueError(
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"GALU was initialized with C={self.channels} but received input with C={int(x.shape[1])}."
+            )
+
+        x_norm = x / self.max_input
+        zero = torch.zeros(1, dtype=x.dtype, device=x.device)
+        alpha = self._reshape(self.alpha, x_norm)
+        beta = self._reshape(self.beta, x_norm)
+        gamma = self._reshape(self.gamma, x_norm)
+        delta = self._reshape(self.delta, x_norm)
+        part_prelu = F.relu(x_norm) + alpha * torch.minimum(x_norm, zero)
+        part_beta = beta * (
+            F.relu(1.0 - torch.abs(x_norm - 1.0))
+            + torch.minimum(torch.abs(x_norm - 3.0) - 1.0, zero)
+        )
+        part_gamma = gamma * (
+            F.relu(0.5 - torch.abs(x_norm - 0.5))
+            + torch.minimum(torch.abs(x_norm - 1.5) - 0.5, zero)
+        )
+        part_delta = delta * (
+            F.relu(0.5 - torch.abs(x_norm - 2.5))
+            + torch.minimum(torch.abs(x_norm - 3.5) - 0.5, zero)
+        )
+        z = part_prelu + part_beta + part_gamma + part_delta
+        return z * self.max_input

homa-0.2.1/src/homa/activations/classes/MELU.py

@@ -0,0 +1,63 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class MELU(nn.Module):
+    def __init__(self, channels: int, max_input: float = 1.0):
+        super().__init__()
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.max_input = float(max_input)
+        shape = (1, self.channels, 1, 1)
+        self.alpha = nn.Parameter(torch.empty(shape))
+        self.beta = nn.Parameter(torch.empty(shape))
+        self.gamma = nn.Parameter(torch.empty(shape))
+        self.delta = nn.Parameter(torch.empty(shape))
+        self.xi = nn.Parameter(torch.empty(shape))
+        self.psi = nn.Parameter(torch.empty(shape))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            for param in (
+                self.alpha,
+                self.beta,
+                self.gamma,
+                self.delta,
+                self.xi,
+                self.psi,
+            ):
+                param.zero_()
+
+    def forward(self, X: torch.Tensor) -> torch.Tensor:
+        if X.dim() != 4:
+            raise ValueError(
+                f"Expected 4D input (N, C, H, W), got {X.dim()}D with shape {tuple(X.shape)}"
+            )
+        if int(X.shape[1]) != self.channels:
+            raise ValueError(
+                f"MELU was initialized with C={self.channels} but received input with C={int(X.shape[1])}."
+            )
+
+        X_norm = X / self.max_input
+        Y = torch.roll(X_norm, shifts=-1, dims=1)
+
+        term1 = F.relu(X_norm)
+        term2 = self.alpha * torch.clamp(X_norm, max=0)
+
+        dist_sq_beta = (X_norm - 2) ** 2 + (Y - 2) ** 2
+        dist_sq_gamma = (X_norm - 1) ** 2 + (Y - 1) ** 2
+        dist_sq_delta = (X_norm - 1) ** 2 + (Y - 3) ** 2
+        dist_sq_xi = (X_norm - 3) ** 2 + (Y - 1) ** 2
+        dist_sq_psi = (X_norm - 3) ** 2 + (Y - 3) ** 2
+
+        term3 = self.beta * torch.sqrt(F.relu(2 - dist_sq_beta))
+        term4 = self.gamma * torch.sqrt(F.relu(1 - dist_sq_gamma))
+        term5 = self.delta * torch.sqrt(F.relu(1 - dist_sq_delta))
+        term6 = self.xi * torch.sqrt(F.relu(1 - dist_sq_xi))
+        term7 = self.psi * torch.sqrt(F.relu(1 - dist_sq_psi))
+
+        Z_norm = term1 + term2 + term3 + term4 + term5 + term6 + term7
+        return Z_norm * self.max_input

homa-0.2.1/src/homa/activations/classes/PDELU.py

@@ -0,0 +1,44 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class PDELU(nn.Module):
+    def __init__(self, channels: int, theta: float = 0.5):
+        super().__init__()
+        if theta == 1.0:
+            raise ValueError("theta cannot be 1.0, as it would cause a division by zero.")
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.theta = float(theta)
+        self._power_val = 1.0 / (1.0 - self.theta)
+        self.alpha = nn.Parameter(torch.empty(self.channels))
+        self.reset_parameters()
+
+    @staticmethod
+    def _reshape(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        return param.view(1, param.shape[0], *([1] * (ref.ndim - 2)))
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            self.alpha.fill_(0.1)
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim < 2:
+            raise ValueError(
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"PDELU was initialized with C={self.channels} but received input with C={int(x.shape[1])}."
+            )
+
+        positive_part = F.relu(x)
+        inner_term = F.relu(1.0 + (1.0 - self.theta) * x)
+        powered_term = torch.pow(inner_term, self._power_val)
+        subtracted_term = powered_term - 1.0
+        zero = torch.zeros(1, dtype=x.dtype, device=x.device)
+        alpha = self._reshape(self.alpha, x)
+        negative_part = alpha * torch.minimum(subtracted_term, zero)
+        return positive_part + negative_part

homa-0.2.1/src/homa/activations/classes/SReLU.py

@@ -0,0 +1,50 @@
+import torch
+from torch import nn
+
+
+class SReLU(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        alpha_init: float = 0.0,
+        beta_init: float = 0.0,
+        gamma_init: float = 1.0,
+        delta_init: float = 1.0,
+    ):
+        super().__init__()
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.alpha_init_val = float(alpha_init)
+        self.beta_init_val = float(beta_init)
+        self.gamma_init_val = float(gamma_init)
+        self.delta_init_val = float(delta_init)
+        shape = (1, self.channels, 1, 1)
+        self.alpha = nn.Parameter(torch.empty(shape))
+        self.beta = nn.Parameter(torch.empty(shape))
+        self.gamma = nn.Parameter(torch.empty(shape))
+        self.delta = nn.Parameter(torch.empty(shape))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            self.alpha.fill_(self.alpha_init_val)
+            self.beta.fill_(self.beta_init_val)
+            self.gamma.fill_(self.gamma_init_val)
+            self.delta.fill_(self.delta_init_val)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 4:
+            raise ValueError(
+                f"Expected 4D input (N, C, H, W), got {x.dim()}D with shape {tuple(x.shape)}"
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"SReLU was initialized with C={self.channels} but received input with C={int(x.shape[1])}."
+            )
+
+        start = self.beta + self.alpha * (x - self.beta)
+        finish = self.delta + self.gamma * (x - self.delta)
+        out = torch.where(x < self.beta, start, x)
+        out = torch.where(x > self.delta, finish, out)
+        return out

homa-0.2.1/src/homa/activations/classes/SmallGALU.py

@@ -0,0 +1,48 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class SmallGALU(nn.Module):
+    def __init__(self, channels: int, max_input: float = 1.0):
+        super().__init__()
+        if max_input <= 0:
+            raise ValueError("max_input must be positive.")
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.max_input = float(max_input)
+        self.alpha = nn.Parameter(torch.empty(self.channels))
+        self.beta = nn.Parameter(torch.empty(self.channels))
+        self.reset_parameters()
+
+    @staticmethod
+    def _reshape(param: torch.Tensor, ref: torch.Tensor) -> torch.Tensor:
+        return param.view(1, param.shape[0], *([1] * (ref.ndim - 2)))
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            self.alpha.zero_()
+            self.beta.zero_()
+
+    def forward(self, x: torch.Tensor):
+        if x.ndim < 2:
+            raise ValueError(
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"SmallGALU was initialized with C={self.channels} but received input with C={int(x.shape[1])}."
+            )
+
+        x_norm = x / self.max_input
+        zero = torch.zeros(1, dtype=x.dtype, device=x.device)
+        alpha = self._reshape(self.alpha, x_norm)
+        beta = self._reshape(self.beta, x_norm)
+        part_prelu = F.relu(x_norm) + alpha * torch.minimum(x_norm, zero)
+        part_beta = beta * (
+            F.relu(1.0 - torch.abs(x_norm - 1.0))
+            + torch.minimum(torch.abs(x_norm - 3.0) - 1.0, zero)
+        )
+        z = part_prelu + part_beta
+        return z * self.max_input

homa-0.2.1/src/homa/activations/classes/WideMELU.py

@@ -0,0 +1,68 @@
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+
+class WideMELU(nn.Module):
+    def __init__(self, channels: int, max_input: float = 1.0):
+        super().__init__()
+        if channels <= 0:
+            raise ValueError(f"channels must be positive, got {channels}.")
+        self.channels = int(channels)
+        self.max_input = float(max_input)
+        shape = (1, self.channels, 1, 1)
+        self.alpha = nn.Parameter(torch.empty(shape))
+        self.beta = nn.Parameter(torch.empty(shape))
+        self.gamma = nn.Parameter(torch.empty(shape))
+        self.delta = nn.Parameter(torch.empty(shape))
+        self.xi = nn.Parameter(torch.empty(shape))
+        self.psi = nn.Parameter(torch.empty(shape))
+        self.theta = nn.Parameter(torch.empty(shape))
+        self.lam = nn.Parameter(torch.empty(shape))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        with torch.no_grad():
+            for param in (
+                self.alpha,
+                self.beta,
+                self.gamma,
+                self.delta,
+                self.xi,
+                self.psi,
+                self.theta,
+                self.lam,
+            ):
+                param.zero_()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if x.dim() != 4:
+            raise ValueError(
+                f"Expected 4D input (N, C, H, W), got {x.dim()}D with shape {tuple(x.shape)}"
+            )
+        if int(x.shape[1]) != self.channels:
+            raise ValueError(
+                f"WideMELU was initialized with C={self.channels} but received input with C={int(x.shape[1])}."
+            )
+
+        X_norm = x / self.max_input
+        Y = torch.roll(X_norm, shifts=-1, dims=1)
+
+        term1 = F.relu(X_norm)
+        term2 = self.alpha * torch.clamp(X_norm, max=0)
+        dist_sq_beta = (X_norm - 2) ** 2 + (Y - 2) ** 2
+        dist_sq_gamma = (X_norm - 1) ** 2 + (Y - 1) ** 2
+        dist_sq_delta = (X_norm - 1) ** 2 + (Y - 3) ** 2
+        dist_sq_xi = (X_norm - 3) ** 2 + (Y - 1) ** 2
+        dist_sq_psi = (X_norm - 3) ** 2 + (Y - 3) ** 2
+        dist_sq_theta = (X_norm - 1) ** 2 + (Y - 2) ** 2
+        dist_sq_lambda = (X_norm - 3) ** 2 + (Y - 2) ** 2
+        term3 = self.beta * torch.sqrt(F.relu(2 - dist_sq_beta))
+        term4 = self.gamma * torch.sqrt(F.relu(1 - dist_sq_gamma))
+        term5 = self.delta * torch.sqrt(F.relu(1 - dist_sq_delta))
+        term6 = self.xi * torch.sqrt(F.relu(1 - dist_sq_xi))
+        term7 = self.psi * torch.sqrt(F.relu(1 - dist_sq_psi))
+        term8 = self.theta * torch.sqrt(F.relu(1 - dist_sq_theta))
+        term9 = self.lam * torch.sqrt(F.relu(1 - dist_sq_lambda))
+        Z_norm = term1 + term2 + term3 + term4 + term5 + term6 + term7 + term8 + term9
+        return Z_norm * self.max_input

homa-0.2.1/src/homa/activations/utils.py

@@ -0,0 +1,94 @@
+from typing import Optional, Sequence, Tuple, Type
+
+import torch
+from torch import nn
+
+
+def negative_part(x):
+    return torch.minimum(x, torch.zeros_like(x))
+
+
+def positive_part(x):
+    return torch.maximum(x, torch.zeros_like(x))
+
+
+def as_channel_parameters(parameter: torch.Tensor, x: torch.Tensor):
+    shape = [1] * x.dim()
+    shape[1] = -1
+    return parameter.view(*shape)
+
+
+def device_compatibility_check(model, x: torch.Tensor):
+    for p in model.parameters():
+        if p.device != x.device or p.dtype != x.dtype:
+            p.data = p.data.to(device=x.device, dtype=x.dtype)
+
+
+def phi_hat(x, a, lam):
+    term_pos = torch.maximum(lam - torch.abs(x - a), torch.zeros_like(x))
+    term_neg = torch.minimum(torch.abs(x - (a + 2 * lam)) - lam, torch.zeros_like(x))
+    return term_pos + term_neg
+
+
+def _channels_from_module(module: nn.Module) -> Optional[int]:
+    for attr in ("channels", "out_channels", "out_features", "num_features"):
+        value = getattr(module, attr, None)
+        if isinstance(value, int) and value > 0:
+            return int(value)
+    weight = getattr(module, "weight", None)
+    if isinstance(weight, torch.Tensor) and weight.ndim >= 1:
+        return int(weight.shape[0])
+    return None
+
+
+def _infer_channels(parent: nn.Module, target_name: str, child: nn.Module) -> int:
+    value = _channels_from_module(child)
+    if value is not None:
+        return value
+
+    children: Sequence[Tuple[str, nn.Module]] = list(parent.named_children())
+    for index, (name, _) in enumerate(children):
+        if name == target_name:
+            break
+    else:
+        raise ValueError(f"Child {target_name} not found in parent {parent.__class__.__name__}.")
+
+    for _, module in reversed(children[:index]):
+        value = _channels_from_module(module)
+        if value is not None:
+            return value
+
+    value = _channels_from_module(parent)
+    if value is not None:
+        return value
+
+    raise ValueError(
+        f"Could not infer channel count for activation {target_name!r} under parent {parent.__class__.__name__}."
+    )
+
+
+def infer_activation_channels(parent: nn.Module, child_name: str, child: nn.Module) -> int:
+    return _infer_channels(parent, child_name, child)
+
+
+def replace_activation(
+    model: nn.Module,
+    activation: Type[nn.Module],
+    needles: Optional[Sequence[Type[nn.Module]]] = None,
+    **activation_kwargs,
+) -> None:
+    if needles is None:
+        needles = (nn.ReLU,)
+    needle_types = tuple(needles)
+
+    for parent in model.modules():
+        for name, child in list(parent.named_children()):
+            if isinstance(child, needle_types):
+                channels = _infer_channels(parent, name, child)
+                try:
+                    new_module = activation(channels=channels, **activation_kwargs)
+                except TypeError as exc:  # pragma: no cover - defensive branch
+                    raise TypeError(
+                        f"{activation.__name__} must accept a `channels` keyword argument."
+                    ) from exc
+                setattr(parent, name, new_module)

homa-0.2.1/src/homa/vision/StochasticClassifier.py

@@ -0,0 +1,50 @@
+import random
+import torch
+from ..activations import (
+    APLU,
+    GALU,
+    SmallGALU,
+    MELU,
+    WideMELU,
+    PDELU,
+    SReLU,
+    infer_activation_channels,
+)
+
+
+class StochasticClassifier:
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self._activation_pool = [
+            APLU,
+            GALU,
+            SmallGALU,
+            MELU,
+            WideMELU,
+            PDELU,
+            SReLU,
+            torch.nn.ReLU,
+            torch.nn.PReLU,
+            torch.nn.LeakyReLU,
+            torch.nn.ELU,
+        ]
+        self._requires_channels = {
+            APLU,
+            GALU,
+            SmallGALU,
+            MELU,
+            WideMELU,
+            PDELU,
+            SReLU,
+        }
+
+    def replace_activations(self, needle: torch.nn.Module) -> None:
+        for parent in self.network.modules():
+            for name, child in list(parent.named_children()):
+                if isinstance(child, needle):
+                    replacement = random.choice(self._activation_pool)
+                    if replacement in self._requires_channels:
+                        channels = infer_activation_channels(parent, name, child)
+                        setattr(parent, name, replacement(channels=channels))
+                    else:
+                        setattr(parent, name, replacement())

{homa-0.2.0 → homa-0.2.1}/src/homa.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: homa
-Version: 0.2.0
+Version: 0.2.1
 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/src/homa/activations/classes/APLU.py

@@ -1,86 +0,0 @@
-import torch
-from torch import nn
-from torch.nn.parameter import Parameter, UninitializedParameter
-import torch.nn.functional as F
-
-
-class APLU(nn.Module):
-    def __init__(self, max_input: float = 1.0):
-        super().__init__()
-        self.max_input = float(max_input)
-        self.alpha = UninitializedParameter()
-        self.beta = UninitializedParameter()
-        self.gamma = UninitializedParameter()
-        self.xi = UninitializedParameter()
-        self.psi = UninitializedParameter()
-        self.mu = UninitializedParameter()
-        self._num_channels = None
-
-    def _initialize_parameters(self, x: torch.Tensor):
-        if x.ndim < 2:
-            raise ValueError(
-                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
-            )
-
-        channels = int(x.shape[1])
-        self._num_channels = channels
-        param_shape = [1] * x.ndim
-        param_shape[1] = channels
-
-        with torch.no_grad():
-            self.alpha = Parameter(
-                torch.zeros(param_shape, dtype=x.dtype, device=x.device)
-            )
-            self.beta = Parameter(
-                torch.zeros(param_shape, dtype=x.dtype, device=x.device)
-            )
-            self.gamma = Parameter(
-                torch.zeros(param_shape, dtype=x.dtype, device=x.device)
-            )
-            self.xi = Parameter(
-                torch.empty(param_shape, dtype=x.dtype, device=x.device).uniform_(
-                    0.0, self.max_input
-                )
-            )
-            self.psi = Parameter(
-                torch.empty(param_shape, dtype=x.dtype, device=x.device).uniform_(
-                    0.0, self.max_input
-                )
-            )
-            self.mu = Parameter(
-                torch.empty(param_shape, dtype=x.dtype, device=x.device).uniform_(
-                    0.0, self.max_input
-                )
-            )
-
-    def reset_parameters(self):
-        if isinstance(self.alpha, UninitializedParameter):
-            return
-
-        with torch.no_grad():
-            self.alpha.zero_()
-            self.beta.zero_()
-            self.gamma.zero_()
-            self.xi.uniform_(0.0, self.max_input)
-            self.psi.uniform_(0.0, self.max_input)
-            self.mu.uniform_(0.0, self.max_input)
-
-    def forward(self, x: torch.Tensor):
-        if isinstance(self.alpha, UninitializedParameter):
-            self._initialize_parameters(x)
-
-        if x.ndim < 2:
-            raise ValueError(
-                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
-            )
-        if self._num_channels is not None and x.shape[1] != self._num_channels:
-            raise RuntimeError(
-                f"APLU was initialized with C={self._num_channels} but got C={x.shape[1]}. "
-                "Create a new APLU for a different channel size."
-            )
-
-        a = F.relu(x)
-        b = self.alpha * F.relu(-x + self.xi)
-        c = self.beta * F.relu(-x + self.psi)
-        d = self.gamma * F.relu(-x + self.mu)
-        return a + b + c + d