Repository: pyro-ppl/numpyro

Issue:
total count as model parameter in multinomial distribution
Here is the issue corresponding to the forum question [5066](https://forum.pyro.ai/t/multinomial-plate-and-total-counts/5066).
Example:
5 Customers are going to buy apples from 4 available category of apples.
- The total number of purchase is negative binomial (with zero purchase being the most common one).
- The probability of choosing apples from each category is drawn from Dirichlet distribution for each customer.
Here is the minimal code:
```python
import jax
import jax.numpy as jnp
import matplotlib.pyplot as plt
import numpyro
import numpyro.distributions as dist
from numpyro.infer import Predictive

seed = 1234
# numpyro seed
rng_key = jax.random.PRNGKey(seed)
rng_trace, rng_prior_pred, rng_posterior_pred = jax.random.split(rng_key, 3)

def gen_data(num_apple_category: int = 4, n_customers: int = 5):
    mu = 1
    alpha = 2
    with numpyro.plate("customers", n_customers):
        # The latent probability for number of buys for each customer is Negative binomial.
        # The lower number of the purchase has more chance.
        num_purchase = numpyro.sample("num_purchase", dist.NegativeBinomial2(mu, alpha))
        desired_probs = numpyro.sample(
            "desired_probs", dist.Dirichlet(0.5 * jnp.ones(num_apple_category))
        )
        numpyro.sample(
            "apples", dist.Multinomial(total_count=num_purchase, probs=desired_probs)
        )


num_samples = 1
# Draw from priors to make the fake data
prior_predictive = Predictive(gen_data, num_samples=num_samples)
prior_predictions = prior_predictive(rng_prior_pred)
```

Error:
```bash
AssertionError: The total count parameter `n` should not be a jax abstract array.
```

Versions:
```bash
jax=='0.4.1'
jaxlib=='0.4.1'
numpyro=='0.10.1'
```


Discussion:

Resolving patch:
diff --git a/numpyro/distributions/discrete.py b/numpyro/distributions/discrete.py
--- a/numpyro/distributions/discrete.py
+++ b/numpyro/distributions/discrete.py
@@ -505,7 +505,9 @@ class MultinomialProbs(Distribution):
         "total_count": constraints.nonnegative_integer,
     }
 
-    def __init__(self, probs, total_count=1, *, validate_args=None):
+    def __init__(
+        self, probs, total_count=1, *, total_count_max=None, validate_args=None
+    ):
         if jnp.ndim(probs) < 1:
             raise ValueError("`probs` parameter must be at least one-dimensional.")
         batch_shape, event_shape = self.infer_shapes(
@@ -513,6 +515,7 @@ def __init__(self, probs, total_count=1, *, validate_args=None):
         )
         self.probs = promote_shapes(probs, shape=batch_shape + jnp.shape(probs)[-1:])[0]
         self.total_count = promote_shapes(total_count, shape=batch_shape)[0]
+        self.total_count_max = total_count_max
         super(MultinomialProbs, self).__init__(
             batch_shape=batch_shape,
             event_shape=event_shape,
@@ -522,7 +525,11 @@ def __init__(self, probs, total_count=1, *, validate_args=None):
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         return multinomial(
-            key, self.probs, self.total_count, shape=sample_shape + self.batch_shape
+            key,
+            self.probs,
+            self.total_count,
+            shape=sample_shape + self.batch_shape,
+            total_count_max=self.total_count_max,
         )
 
     @validate_sample
@@ -562,7 +569,9 @@ class MultinomialLogits(Distribution):
         "total_count": constraints.nonnegative_integer,
     }
 
-    def __init__(self, logits, total_count=1, *, validate_args=None):
+    def __init__(
+        self, logits, total_count=1, *, total_count_max=None, validate_args=None
+    ):
         if jnp.ndim(logits) < 1:
             raise ValueError("`logits` parameter must be at least one-dimensional.")
         batch_shape, event_shape = self.infer_shapes(
@@ -572,6 +581,7 @@ def __init__(self, logits, total_count=1, *, validate_args=None):
             logits, shape=batch_shape + jnp.shape(logits)[-1:]
         )[0]
         self.total_count = promote_shapes(total_count, shape=batch_shape)[0]
+        self.total_count_max = total_count_max
         super(MultinomialLogits, self).__init__(
             batch_shape=batch_shape,
             event_shape=event_shape,
@@ -581,7 +591,11 @@ def __init__(self, logits, total_count=1, *, validate_args=None):
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
         return multinomial(
-            key, self.probs, self.total_count, shape=sample_shape + self.batch_shape
+            key,
+            self.probs,
+            self.total_count,
+            shape=sample_shape + self.batch_shape,
+            total_count_max=self.total_count_max,
         )
 
     @validate_sample
@@ -618,11 +632,32 @@ def infer_shapes(logits, total_count):
         return batch_shape, event_shape
 
 
-def Multinomial(total_count=1, probs=None, logits=None, *, validate_args=None):
+def Multinomial(
+    total_count=1, probs=None, logits=None, *, total_count_max=None, validate_args=None
+):
+    """Multinomial distribution.
+
+    :param total_count: number of trials. If this is a JAX array,
+        it is required to specify `total_count_max`.
+    :param probs: event probabilities
+    :param logits: event log probabilities
+    :param int total_count_max: the maximum number of trials,
+        i.e. `max(total_count)`
+    """
     if probs is not None:
-        return MultinomialProbs(probs, total_count, validate_args=validate_args)
+        return MultinomialProbs(
+            probs,
+            total_count,
+            total_count_max=total_count_max,
+            validate_args=validate_args,
+        )
     elif logits is not None:
-        return MultinomialLogits(logits, total_count, validate_args=validate_args)
+        return MultinomialLogits(
+            logits,
+            total_count,
+            total_count_max=total_count_max,
+            validate_args=validate_args,
+        )
     else:
         raise ValueError("One of `probs` or `logits` must be specified.")
 
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -260,11 +260,15 @@ def _multinomial(key, p, n, n_max, shape=()):
     return jnp.reshape(samples_2D, shape + p.shape[-1:]) - excess
 
 
-def multinomial(key, p, n, shape=()):
-    assert not isinstance(
-        n, jax.core.Tracer
-    ), "The total count parameter `n` should not be a jax abstract array."
-    n_max = int(np.max(jax.device_get(n)))
+def multinomial(key, p, n, shape=(), total_count_max=None):
+    if total_count_max is None:
+        if isinstance(n, jax.core.Tracer):
+            raise ValueError(
+                "Please specify total_count_max in Multinomial distribution."
+            )
+        n_max = int(np.max(jax.device_get(n)))
+    else:
+        n_max = total_count_max
     return _multinomial(key, p, n, n_max, shape)
 
 


Test patch:
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -187,6 +187,8 @@ def model(data):
 
 @pytest.mark.parametrize("kernel_cls", [HMC, NUTS, SA, BarkerMH])
 def test_beta_bernoulli_x64(kernel_cls):
+    if kernel_cls is SA and "CI" in os.environ and "JAX_ENABLE_X64" in os.environ:
+        pytest.skip("The test is flaky on CI x64.")
     num_warmup, num_samples = (100000, 100000) if kernel_cls is SA else (500, 20000)
 
     def model(data):
@@ -318,6 +320,8 @@ def model(data):
 
 @pytest.mark.parametrize("with_logits", ["True", "False"])
 def test_binomial_stable_x64(with_logits):
+    if "CI" in os.environ and "JAX_ENABLE_X64" in os.environ:
+        pytest.skip("The test is flaky on CI x64.")
     # Ref: https://github.com/pyro-ppl/pyro/issues/1706
     num_warmup, num_samples = 200, 200
 
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -2743,3 +2743,18 @@ def sample_binomial_withp0(key):
         return dist.Binomial(total_count=n, probs=0).sample(key)
 
     jax.vmap(sample_binomial_withp0)(random.split(random.PRNGKey(0), 1))
+
+
+def test_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.Multinomial(total_count, probs=probs, total_count_max=10).sample(
+            key
+        )
+
+    x = dist.Multinomial(10, probs).sample(key)
+    y = jax.jit(f)(x)
+    assert_allclose(x, y, rtol=1e-6)