distribution output

pts/modules/distribution_output.py

@@ -4,12 +4,15 @@ from typing import Callable, Dict, Optional, Tuple
 import numpy as np
 import torch
 import torch.nn as nn
+from torch.distributions import Distribution, TransformedDistribution, AffineTransform
+
+from .lambda_layer import LambdaLayer
 
 
 class ArgProj(nn.Module):
     def __init__(
         self,
-        in_features,
+        in_features: int,
         args_dim: Dict[str, int],
         domain_map: Callable[..., Tuple[torch.Tensor]],
         dtype: np.dtype = np.float32,
@@ -31,4 +34,42 @@ class ArgProj(nn.Module):
 
 
 class Output(ABC):
-    pass
+    in_features: int
+    args_dim: Dict[str, int]
+    _dtype: np.dtype = np.float32
+
+    @property
+    def dtype(self):
+        return self._dtype
+
+    @dtype.setter
+    def dtype(self, dtype: np.dtype):
+        self._dtype = dtype
+
+    def get_args_proj(self, prefix: Optional[str] = None) -> ArgProj:
+        return ArgProj(
+            in_features=self.in_features,
+            args_dim=self.args_dim,
+            domain_map=LambdaLayer(self.domain_map),
+            prefix=prefix,
+            dtype=self.dtype,
+        )
+
+    @abstractmethod
+    def domain_map(self, *args: torch.Tensor):
+        pass
+
+
+class DistributionOutput(Output):
+
+    distr_cls: type
+
+    def distribution(
+        self, distr_args, scale: Optional[torch.Tensor] = None
+    ) -> Distribution:
+
+        if scale is None:
+            return self.distr_cls(*distr_args)
+        else:
+            distr = self.distr_cls(*distr_args)
+            return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)])