total count as model parameter in multinomial distribution

[mahdikooshkbaghi]

Here is the issue corresponding to the forum question 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:

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:

AssertionError: The total count parameter `n` should not be a jax abstract array.

Versions:

jax=='0.4.1'
jaxlib=='0.4.1'
numpyro=='0.10.1'