fix(distributions.conjugate): support total_count_max in DirichletMultinomial distribution (#2016)

* feat(distributions.conjugate): support total_count_max in DirichletMultinomial

* fix(distributions.conjugate): add total_count_max to pytree_aux_fields

* test: mark neutra example as xfail

* fix: rm pngs and graph, gitignore for posterity

---------

Co-authored-by: Brendan Cooley <[email protected]>

Author: brendancooley

.gitignore

@@ -19,6 +19,7 @@ build
 
 # data files
 *.pdf
+*.png
 numpyro/examples/.data
 examples/.results
 examples/*.pdf

numpyro/distributions/conjugate.py

@@ -107,16 +107,19 @@ class DirichletMultinomial(Distribution):
     :param numpy.ndarray concentration: concentration parameter (alpha) for the
         Dirichlet distribution.
     :param numpy.ndarray total_count: number of Categorical trials.
+    :param int total_count_max: the maximum number of trials, i.e. `max(total_count)`
     """
 
     arg_constraints = {
         "concentration": constraints.independent(constraints.positive, 1),
         "total_count": constraints.nonnegative_integer,
     }
     pytree_data_fields = ("concentration", "_dirichlet")
-    pytree_aux_fields = ("total_count",)
+    pytree_aux_fields = ("total_count", "total_count_max")
 
-    def __init__(self, concentration, total_count=1, *, validate_args=None):
+    def __init__(
+        self, concentration, total_count=1, *, total_count_max=None, validate_args=None
+    ):
         if jnp.ndim(concentration) < 1:
             raise ValueError(
                 "`concentration` parameter must be at least one-dimensional."
@@ -128,6 +131,7 @@ def __init__(self, concentration, total_count=1, *, validate_args=None):
         concentration_shape = batch_shape + jnp.shape(concentration)[-1:]
         (self.concentration,) = promote_shapes(concentration, shape=concentration_shape)
         (self.total_count,) = promote_shapes(total_count, shape=batch_shape)
+        self.total_count_max = total_count_max
         concentration = jnp.broadcast_to(self.concentration, concentration_shape)
         self._dirichlet = Dirichlet(concentration)
         super().__init__(
@@ -140,9 +144,11 @@ def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         key_dirichlet, key_multinom = random.split(key)
         probs = self._dirichlet.sample(key_dirichlet, sample_shape)
-        return MultinomialProbs(total_count=self.total_count, probs=probs).sample(
-            key_multinom
-        )
+        return MultinomialProbs(
+            total_count=self.total_count,
+            probs=probs,
+            total_count_max=self.total_count_max,
+        ).sample(key_multinom)
 
     @validate_sample
     def log_prob(self, value):

test/test_distributions.py

@@ -3416,6 +3416,23 @@ def f(x):
     assert_allclose(x, y, rtol=1e-6)
 
 
+def test_dirichlet_multinomial_abstract_total_count():
+    probs = jnp.array([0.2, 0.5, 0.3])
+    key = random.PRNGKey(0)
+
+    def f(x):
+        total_count = x.sum(-1)
+        return dist.DirichletMultinomial(
+            concentration=probs,
+            total_count=total_count,
+            total_count_max=10,  # fails on 0.18.0
+        ).sample(key)
+
+    x = dist.DirichletMultinomial(concentration=probs, total_count=10).sample(key)
+    y = jax.jit(f)(x)
+    assert_allclose(x, y, rtol=1e-6)
+
+
 def test_normal_log_cdf():
     # test if log_cdf method agrees with jax.scipy.stats.norm.logcdf
     # and if exp(log_cdf) agrees with cdf

test/test_examples.py

@@ -49,7 +49,13 @@
     "hsgp.py --num-samples 10 --num-warmup 10 --num-chains 2",
     "minipyro.py",
     "mortality.py --num-samples 10 --num-warmup 10 --num-chains 2",
-    "neutra.py --num-samples 100 --num-warmup 100",
+    pytest.param(
+        "neutra.py --num-samples 100 --num-warmup 100",
+        marks=pytest.mark.skipif(
+            "CI" in os.environ,
+            reason="This example fails on the CI runner with message 'died with <Signals.SIGSEGV: 11>.'",
+        ),
+    ),
     "ode.py --num-samples 100 --num-warmup 100 --num-chains 1",
     pytest.param(
         "prodlda.py --num-steps 10 --hidden 10 --nn-framework flax",