An alternative implementation of general finite mixture models (pyro-…

…ppl#1458) * initial draft of MixtureGeneral distribution * fixing lint adding some tests * adding some docstrings and tests for general mixture * removing hanging comment * fixing docstring parameter name
DanWaxman · Aug 8, 2022 · be1eb6b · be1eb6b
1 parent 0607a92
commit be1eb6b
Show file tree

Hide file tree

Showing 5 changed files with 452 additions and 142 deletions.
diff --git a/docs/source/distributions.rst b/docs/source/distributions.rst
@@ -532,6 +532,10 @@ ZeroInflatedNegativeBinomial2
 Mixture Distributions
 ---------------------
 
+Mixture
+^^^^^^^
+.. autofunction:: numpyro.distributions.mixtures.Mixture
+
 MixtureSameFamily
 ^^^^^^^^^^^^^^^^^
 .. autoclass:: numpyro.distributions.mixtures.MixtureSameFamily
@@ -540,6 +544,13 @@ MixtureSameFamily
     :show-inheritance:
     :member-order: bysource
 
+MixtureGeneral
+^^^^^^^^^^^^^^
+.. autoclass:: numpyro.distributions.mixtures.MixtureGeneral
+    :members:
+    :undoc-members:
+    :show-inheritance:
+    :member-order: bysource
 
 Directional Distributions
 -------------------------

diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
@@ -84,7 +84,7 @@
     Unit,
 )
 from numpyro.distributions.kl import kl_divergence
-from numpyro.distributions.mixtures import MixtureSameFamily
+from numpyro.distributions.mixtures import Mixture, MixtureGeneral, MixtureSameFamily
 from numpyro.distributions.transforms import biject_to
 from numpyro.distributions.truncated import (
     LeftTruncatedDistribution,
@@ -147,7 +147,9 @@
     "Logistic",
     "LogNormal",
     "MaskedDistribution",
+    "Mixture",
     "MixtureSameFamily",
+    "MixtureGeneral",
     "Multinomial",
     "MultinomialLogits",
     "MultinomialProbs",