initial modules

pts/modules/distribution_output.py

@@ -0,0 +1,32 @@
+from typing import Callable, Dict, Optional, Tuple
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+
+class ArgProj(nn.Module):
+    def __init__(
+            self,
+            in_features,
+            args_dim: Dict[str, int],
+            domain_map: Callable[..., Tuple[torch.Tensor]],
+            dtype: np.dtype = np.float32,
+            prefix: Optional[str] = None,
+            **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.args_dim = args_dim
+        self.dtype = dtype
+        self.proj = nn.ModuleList([
+            nn.Linear(in_features, dim)
+            for dim in args_dim.values()])
+        self.domain_map = domain_map
+    
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor]:
+        params_unbounded = [proj(x) for proj in self.proj]
+
+        return self.domain_map(*params_unbounded)
+
+

pts/modules/lambda.py

@@ -0,0 +1,11 @@
+import torch
+import torch.nn as nn
+
+class Lambda(nn.Module):
+    def __init(self, function):
+        super().__init__()
+        func_dict = {sym: getattr(sym, function), Tensor: getattr(torch.Tensor, function)}
+        self._func = lambda *args: func_dict(*args)
+    
+    def forward(self, x, *args):
+        return self._func(x, *args)