An alternative implementation of general finite mixture models (#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

Author: dfm

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
 -------------------------

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",