fix(gh-2123): Validate component distributions of MixtureSameFamily have parameter-independent support (#2127)

fixes #2123

Author: Qazalbash

numpyro/distributions/mixtures.py

@@ -222,6 +222,13 @@ def __init__(
         *,
         validate_args: Optional[bool] = None,
     ):
+        assert isinstance(
+            component_distribution.support, constraints.ParameterFreeConstraint
+        ), (
+            f"Invalid component distribution: {type(component_distribution).__name__}. "
+            "The mixture components must have a support that does not depend on their parameters "
+            f"(expected ParameterFreeConstraint, but found {component_distribution.support})."
+        )
         _check_mixing_distribution(mixing_distribution)
         mixture_size = mixing_distribution.probs.shape[-1]
         if not isinstance(component_distribution, Distribution):

test/test_distributions_mixture.py

@@ -1,6 +1,8 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
+
+import numpy as np
 import pytest
 
 import jax
@@ -175,3 +177,31 @@ def _test_mixture(mixing_distribution, component_distribution):
     if mixture.event_shape == ():
         cdf = mixture.cdf(samples)
         assert cdf.shape == (*sample_shape, *mixture.shape())
+
+
+@pytest.mark.parametrize(
+    "component_dist",
+    [
+        dist.Uniform(low=np.array([0.0, 5.0]), high=np.array([1.0, 10.0])),
+        dist.TruncatedNormal(loc=np.zeros(2), scale=np.ones(2), low=0.0, high=1.0),
+    ],
+)
+def test_mixture_rejects_parameter_dependent_components(component_dist):
+    mixing_dist = dist.Categorical(probs=np.array([0.5, 0.5]))
+    with pytest.raises(
+        AssertionError, match="expected ParameterFreeConstraint, but found "
+    ):
+        dist.MixtureSameFamily(mixing_dist, component_dist)
+
+
+@pytest.mark.parametrize(
+    "component_dist",
+    [
+        dist.Normal(np.zeros(2), np.ones(2)),
+        dist.Exponential(np.ones(2)),
+        dist.Bernoulli(probs=np.array([0.5, 0.5])),
+    ],
+)
+def test_mixture_accepts_parameter_free_components(component_dist):
+    mixing_dist = dist.Categorical(probs=np.array([0.3, 0.7]))
+    dist.MixtureSameFamily(mixing_dist, component_dist)