add Multinomial distribution

pixyz/distributions/__init__.py

@@ -5,6 +5,7 @@
     FactorizedBernoulli,
     Categorical,
     RelaxedCategorical,
+    Multinomial,
     Dirichlet,
     Beta,
     Laplace,
@@ -45,6 +46,7 @@
     'FactorizedBernoulli',
     'Categorical',
     'RelaxedCategorical',
+    'Multinomial',
     'Dirichlet',
     'Beta',
     'Laplace',

pixyz/distributions/exponential_distributions.py

@@ -5,6 +5,7 @@
     import RelaxedOneHotCategorical as RelaxedOneHotCategoricalTorch
 from torch.distributions.one_hot_categorical\
     import OneHotCategorical as CategoricalTorch
+from torch.distributions import Multinomial as MultinomialTorch
 from torch.distributions import Dirichlet as DirichletTorch
 from torch.distributions import Beta as BetaTorch
 from torch.distributions import Laplace as LaplaceTorch
@@ -167,6 +168,21 @@ def log_likelihood(self, x):
         log_like = self._get_log_like(x)
         return sum_samples(log_like)
 
+class Multinomial(DistributionBase):
+    """
+    Multinomial distribution parameterized by :attr:`total_count` and :attr:`probs`.
+    """
+
+    def __init__(self, cond_var=[], var=["x"], name="p", dim=None, **kwargs):
+        self.params_keys = ["total_count", "probs"]
+        self.DistributionTorch = MultinomialTorch
+
+        super().__init__(cond_var=cond_var, var=var, name=name, dim=dim, **kwargs)
+
+    @property
+    def distribution_name(self):
+        return "Multinomial"
+
 
 class Dirichlet(DistributionBase):
     """
@@ -230,4 +246,3 @@ def __init__(self, cond_var=[], var=["x"], name="p", dim=None, **kwargs):
     @property
     def distribution_name(self):
         return "Gamma"
-