homa 0.1.0__py3-none-any.whl → 0.2.0__py3-none-any.whl

homa/activations/classes/APLU.py

@@ -1,48 +1,86 @@
 import torch
+from torch import nn
+from torch.nn.parameter import Parameter, UninitializedParameter
+import torch.nn.functional as F
 
 
-class APLU(torch.nn.Module):
+class APLU(nn.Module):
     def __init__(self, max_input: float = 1.0):
-        super(APLU, self).__init__()
-        self.max_input = max_input
-        self.alpha = None
-        self.beta = None
-        self.gamma = None
-        self.xi = None
-        self.psi = None
-        self.mu = None
+        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):
+    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 {x.shape}"
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
             )
 
-        num_channels = x.shape[1]
-        self._num_channels = num_channels
-
+        channels = int(x.shape[1])
+        self._num_channels = channels
         param_shape = [1] * x.ndim
-        param_shape[1] = num_channels
+        param_shape[1] = channels
 
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(param_shape))
+        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
+                )
+            )
 
-        self.xi = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
-        self.psi = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
-        self.mu = torch.nn.Parameter(self.max_input * torch.rand(param_shape))
+    def reset_parameters(self):
+        if isinstance(self.alpha, UninitializedParameter):
+            return
 
-    def forward(self, x):
-        if self.alpha is None:
-            self._initialize_parameters(x)
+        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)
 
-        a = torch.relu(x)
+    def forward(self, x: torch.Tensor):
+        if isinstance(self.alpha, UninitializedParameter):
+            self._initialize_parameters(x)
 
-        # following are called hinges
-        b = self.alpha * torch.relu(-x + self.xi)
-        c = self.beta * torch.relu(-x + self.psi)
-        d = self.gamma * torch.relu(-x + self.mu)
-        z = a + b + c + d
+        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."
+            )
 
-        return z
+        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

homa/activations/classes/GALU.py

@@ -1,51 +1,67 @@
 import torch
+from torch import nn
+from torch.nn.parameter import Parameter, UninitializedParameter
+import torch.nn.functional as F
 
 
-class GALU(torch.nn.Module):
+class GALU(nn.Module):
     def __init__(self, max_input: float = 1.0):
-        super(GALU, self).__init__()
+        super().__init__()
         if max_input <= 0:
             raise ValueError("max_input must be positive.")
-        self.max_input = max_input
-        self.alpha = None
-        self.beta = None
-        self.gamma = None
-        self.delta = None
-        self._num_channels = None
+        self.max_input = float(max_input)
+        self.alpha: torch.Tensor = UninitializedParameter()
+        self.beta: torch.Tensor = UninitializedParameter()
+        self.gamma: torch.Tensor = UninitializedParameter()
+        self.delta: torch.Tensor = UninitializedParameter()
 
-    def _initialize_parameters(self, x):
+    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 {x.shape}"
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
             )
-
-        num_channels = x.shape[1]
-        self._num_channels = num_channels
         param_shape = [1] * x.ndim
-        param_shape[1] = num_channels
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(param_shape))
-        self.delta = torch.nn.Parameter(torch.zeros(param_shape))
+        param_shape[1] = int(x.shape[1])
+        zeros = torch.zeros(param_shape, dtype=x.dtype, device=x.device)
+        with torch.no_grad():
+            for name in ("alpha", "beta", "gamma", "delta"):
+                setattr(self, name, Parameter(zeros.clone()))
+
+    def reset_parameters(self):
+        for name in ("alpha", "beta", "gamma", "delta"):
+            p = getattr(self, name)
+            if not isinstance(p, UninitializedParameter):
+                with torch.no_grad():
+                    p.zero_()
 
-    def forward(self, x):
-        if self.alpha is None:
+    def forward(self, x: torch.Tensor):
+        if isinstance(self.alpha, UninitializedParameter):
             self._initialize_parameters(x)
 
-        zero = torch.tensor(0.0, device=x.device, dtype=x.dtype)
+        if x.ndim < 2:
+            raise ValueError(
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
+            )
+        if not isinstance(self.alpha, UninitializedParameter) and x.shape[1] != self.alpha.shape[1]:
+            raise RuntimeError(
+                f"GALU was initialized with C={self.alpha.shape[1]} but got C={x.shape[1]}. "
+                "Create a new GALU for a different channel size."
+            )
+
         x_norm = x / self.max_input
-        part_prelu = torch.relu(x_norm) + self.alpha * torch.min(x_norm, zero)
+        zero = x.new_zeros(1)
+        part_prelu = F.relu(x_norm) + self.alpha * torch.minimum(x_norm, zero)
         part_beta = self.beta * (
-            torch.relu(1.0 - torch.abs(x_norm - 1.0))
-            + torch.min(torch.abs(x_norm - 3.0) - 1.0, zero)
+            F.relu(1.0 - torch.abs(x_norm - 1.0))
+            + torch.minimum(torch.abs(x_norm - 3.0) - 1.0, zero)
         )
         part_gamma = self.gamma * (
-            torch.relu(0.5 - torch.abs(x_norm - 0.5))
-            + torch.min(torch.abs(x_norm - 1.5) - 0.5, zero)
+            F.relu(0.5 - torch.abs(x_norm - 0.5))
+            + torch.minimum(torch.abs(x_norm - 1.5) - 0.5, zero)
         )
         part_delta = self.delta * (
-            torch.relu(0.5 - torch.abs(x_norm - 2.5))
-            + torch.min(torch.abs(x_norm - 3.5) - 0.5, zero)
+            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/activations/classes/MELU.py

@@ -1,50 +1,70 @@
 import torch
+from torch import nn
+import torch.nn.functional as F
 
 
-class MELU(torch.nn.Module):
+class MELU(nn.Module):
     def __init__(self, maxInput: float = 1.0):
         super().__init__()
         self.maxInput = float(maxInput)
-        self.alpha = None
-        self.beta = None
-        self.gamma = None
-        self.delta = None
-        self.xi = None
-        self.psi = None
-        self._initialized = False
-
-    def _initialize_parameters(self, X: torch.Tensor):
-        if X.dim() != 4:
+        self._num_channels = None
+        self.register_parameter("alpha", None)
+        self.register_parameter("beta", None)
+        self.register_parameter("gamma", None)
+        self.register_parameter("delta", None)
+        self.register_parameter("xi", None)
+        self.register_parameter("psi", None)
+
+    def _ensure_parameters(self, x: torch.Tensor):
+        if x.dim() != 4:
             raise ValueError(
-                f"Expected 4D input (B, C, H, W), but got {X.dim()}D input."
+                f"Expected 4D input (N, C, H, W), got {x.dim()}D with shape {tuple(x.shape)}"
+            )
+        c = int(x.shape[1])
+        if self._num_channels is None:
+            self._num_channels = c
+        elif c != self._num_channels:
+            raise RuntimeError(
+                f"MELU was initialized with C={self._num_channels} but got C={c}. "
+                "Create a new MELU for a different channel size."
             )
-        num_channels = X.shape[1]
-        shape = (1, num_channels, 1, 1)
-        self.alpha = torch.nn.Parameter(torch.zeros(shape))
-        self.beta = torch.nn.Parameter(torch.zeros(shape))
-        self.gamma = torch.nn.Parameter(torch.zeros(shape))
-        self.delta = torch.nn.Parameter(torch.zeros(shape))
-        self.xi = torch.nn.Parameter(torch.zeros(shape))
-        self.psi = torch.nn.Parameter(torch.zeros(shape))
-        self._initialized = True
+
+        if self.alpha is None:
+            shape = (1, c, 1, 1)
+            device, dtype = x.device, x.dtype
+            for name in ("alpha", "beta", "gamma", "delta", "xi", "psi"):
+                setattr(
+                    self,
+                    name,
+                    nn.Parameter(torch.zeros(shape, dtype=dtype, device=device)),
+                )
+
+    def reset_parameters(self):
+        for p in (self.alpha, self.beta, self.gamma, self.delta, self.xi, self.psi):
+            if p is not None:
+                with torch.no_grad():
+                    p.zero_()
 
     def forward(self, X: torch.Tensor) -> torch.Tensor:
-        if not self._initialized:
-            self._initialize_parameters(X)
+        self._ensure_parameters(X)
+
         X_norm = X / self.maxInput
         Y = torch.roll(X_norm, shifts=-1, dims=1)
-        term1 = torch.relu(X_norm)
+
+        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(torch.relu(2 - dist_sq_beta))
-        term4 = self.gamma * torch.sqrt(torch.relu(1 - dist_sq_gamma))
-        term5 = self.delta * torch.sqrt(torch.relu(1 - dist_sq_delta))
-        term6 = self.xi * torch.sqrt(torch.relu(1 - dist_sq_xi))
-        term7 = self.psi * torch.sqrt(torch.relu(1 - dist_sq_psi))
+
+        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
-        Z = Z_norm * self.maxInput
-        return Z
+        return Z_norm * self.maxInput

homa/activations/classes/PDELU.py

@@ -1,39 +1,54 @@
 import torch
+from torch import nn
+import torch.nn.functional as F
 
 
-class PDELU(torch.nn.Module):
+class PDELU(nn.Module):
     def __init__(self, theta: float = 0.5):
-        super(PDELU, self).__init__()
+        super().__init__()
         if theta == 1.0:
             raise ValueError(
                 "theta cannot be 1.0, as it would cause a division by zero."
             )
-        self.theta = theta
+        self.theta = float(theta)
         self._power_val = 1.0 / (1.0 - self.theta)
-        self.alpha = torch.nn.UninitializedParameter()
+        self.register_parameter("alpha", None)
         self._num_channels = None
 
-    def _initialize_parameters(self, x: torch.Tensor):
+    def _ensure_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 {x.shape}"
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
             )
 
-        num_channels = x.shape[1]
-        self._num_channels = num_channels
-        param_shape = [1] * x.ndim
-        param_shape[1] = num_channels
-        init_tensor = torch.zeros(param_shape) + 0.1
-        self.alpha = torch.nn.Parameter(init_tensor)
+        c = int(x.shape[1])
+        if self._num_channels is None:
+            self._num_channels = c
+        elif c != self._num_channels:
+            raise RuntimeError(
+                f"PDELU was initialized with C={self._num_channels} but got C={c}. "
+                "Create a new PDELU for a different channel size."
+            )
 
-    def forward(self, x: torch.Tensor):
         if self.alpha is None:
-            self._initialize_parameters(x)
+            param_shape = [1] * x.ndim
+            param_shape[1] = c
+            self.alpha = nn.Parameter(
+                torch.full(param_shape, 0.1, dtype=x.dtype, device=x.device)
+            )
+
+    def reset_parameters(self):
+        if self.alpha is not None:
+            with torch.no_grad():
+                self.alpha.fill_(0.1)
+
+    def forward(self, x: torch.Tensor):
+        self._ensure_parameters(x)
 
-        zero = torch.tensor(0.0, device=x.device, dtype=x.dtype)
-        positive_part = torch.relu(x)
-        inner_term = torch.relu(1.0 + (1.0 - self.theta) * x)
+        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
-        negative_part = self.alpha * torch.min(subtracted_term, zero)
+        zero = torch.zeros(1, dtype=x.dtype, device=x.device)
+        negative_part = self.alpha * torch.minimum(subtracted_term, zero)
         return positive_part + negative_part

homa/activations/classes/SReLU.py

@@ -1,7 +1,8 @@
 import torch
+from torch import nn
 
 
-class SReLU(torch.nn.Module):
+class SReLU(nn.Module):
     def __init__(
         self,
         alpha_init: float = 0.0,
@@ -10,38 +11,57 @@ class SReLU(torch.nn.Module):
         delta_init: float = 1.0,
     ):
         super().__init__()
-        self.alpha_init_val = alpha_init
-        self.beta_init_val = beta_init
-        self.gamma_init_val = gamma_init
-        self.delta_init_val = delta_init
-        self.alpha = torch.nn.UninitializedParameter()
-        self.beta = torch.nn.UninitializedParameter()
-        self.gamma = torch.nn.UninitializedParameter()
-        self.delta = torch.nn.UninitializedParameter()
+        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)
+        self._num_channels = None
+        self.register_parameter("alpha", None)
+        self.register_parameter("beta", None)
+        self.register_parameter("gamma", None)
+        self.register_parameter("delta", None)
 
-    def _initialize_parameters(self, x: torch.Tensor):
-        if isinstance(self.alpha, torch.nn.UninitializedParameter):
-            if x.dim() < 2:
-                raise ValueError(
-                    f"Input tensor must have at least 2 dimensions (N, C), but got {x.dim()}"
-                )
+    def _ensure_parameters(self, x: 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)}"
+            )
+        c = int(x.shape[1])
+        if self._num_channels is None:
+            self._num_channels = c
+        elif c != self._num_channels:
+            raise RuntimeError(
+                f"SReLU was initialized with C={self._num_channels} but got C={c}. "
+                "Create a new SReLU for different channel sizes."
+            )
 
-            num_channels = x.shape[1]
-            param_shape = [1] * x.dim()
-            param_shape[1] = num_channels
-            self.alpha = torch.nn.Parameter(
-                torch.full(param_shape, self.alpha_init_val)
+        if self.alpha is None:
+            shape = (1, c, 1, 1)
+            device, dtype = x.device, x.dtype
+            self.alpha = nn.Parameter(
+                torch.full(shape, self.alpha_init_val, dtype=dtype, device=device)
+            )
+            self.beta = nn.Parameter(
+                torch.full(shape, self.beta_init_val, dtype=dtype, device=device)
             )
-            self.beta = torch.nn.Parameter(torch.full(param_shape, self.beta_init_val))
-            self.gamma = torch.nn.Parameter(
-                torch.full(param_shape, self.gamma_init_val)
+            self.gamma = nn.Parameter(
+                torch.full(shape, self.gamma_init_val, dtype=dtype, device=device)
             )
-            self.delta = torch.nn.Parameter(
-                torch.full(param_shape, self.delta_init_val)
+            self.delta = nn.Parameter(
+                torch.full(shape, self.delta_init_val, dtype=dtype, device=device)
             )
 
+    def reset_parameters(self):
+        if self.alpha is not None:
+            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:
-        self._initialize_parameters(x)
+        self._ensure_parameters(x)
+
         start = self.beta + self.alpha * (x - self.beta)
         finish = self.delta + self.gamma * (x - self.delta)
         out = torch.where(x < self.beta, start, x)

homa/activations/classes/SmallGALU.py

@@ -1,39 +1,58 @@
 import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+import torch.nn.functional as F
 
 
-class SmallGALU(torch.nn.Module):
+class SmallGALU(nn.Module):
     def __init__(self, max_input: float = 1.0):
-        super(SmallGALU, self).__init__()
+        super().__init__()
         if max_input <= 0:
             raise ValueError("max_input must be positive.")
-        self.max_input = max_input
-        self.alpha = None
-        self.beta = None
+        self.max_input = float(max_input)
+        self.register_parameter("alpha", None)
+        self.register_parameter("beta", None)
         self._num_channels = None
 
-    def _initialize_parameters(self, x):
+    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 {x.shape}"
+                f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
             )
-
-        num_channels = x.shape[1]
-        self._num_channels = num_channels
+        self._num_channels = int(x.shape[1])
         param_shape = [1] * x.ndim
-        param_shape[1] = num_channels
-        self.alpha = torch.nn.Parameter(torch.zeros(param_shape))
-        self.beta = torch.nn.Parameter(torch.zeros(param_shape))
+        param_shape[1] = self._num_channels
+        device = x.device
+        dtype = x.dtype
+        self.alpha = Parameter(torch.zeros(param_shape, dtype=dtype, device=device))
+        self.beta = Parameter(torch.zeros(param_shape, dtype=dtype, device=device))
+
+    def reset_parameters(self):
+        if self.alpha is not None:
+            with torch.no_grad():
+                self.alpha.zero_()
+                self.beta.zero_()
 
-    def forward(self, x):
+    def forward(self, x: torch.Tensor):
         if self.alpha is None:
             self._initialize_parameters(x)
+        else:
+            if x.ndim < 2:
+                raise ValueError(
+                    f"Input tensor must have at least 2 dimensions (N, C), but got shape {tuple(x.shape)}"
+                )
+            if x.shape[1] != self._num_channels:
+                raise RuntimeError(
+                    f"SmallGALU was initialized with C={self._num_channels} but got C={x.shape[1]}. "
+                    "Create a new SmallGALU for a different channel size."
+                )
 
-        zero = torch.tensor(0.0, device=x.device, dtype=x.dtype)
         x_norm = x / self.max_input
-        part_prelu = torch.relu(x_norm) + self.alpha * torch.min(x_norm, zero)
+        zero = torch.zeros(1, dtype=x.dtype, device=x.device)
+        part_prelu = F.relu(x_norm) + self.alpha * torch.minimum(x_norm, zero)
         part_beta = self.beta * (
-            torch.relu(1.0 - torch.abs(x_norm - 1.0))
-            + torch.min(torch.abs(x_norm - 3.0) - 1.0, zero)
+            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