Hi! I have an LDA-like model, in which I want to sample `z`

and `w`

with multinomial distributions. We know that in LDA, word is assumed to belong to a topic `z`

. Therefore, In my model. I sampled `z`

like this:

```
z_b_t = numpyro.sample(f"z_b_{t}",
dist.Multinomial(browsed_cards_t,
theta_b_t,
total_count_max=max_browsed_cards_num),
```

The shape of `z_b_t`

is `(batch_size, topic_num)`

, and the elements in `topic_num`

dim represents the occurrences of each topic. I hope to implement the following procedure using multinomial distribution

```
# Choose topic-word distribution first, the topic-word distribution is represented as `phi` with shape is (topic_num, vocabulary_size)
p_z = Vindex(phi)[z_b_t]
# Sample word
word = numpyro.sample("word", dist.Categorical(p_z), obs=obs_word)
```

Because in the procedure of sampling word, we need to choose topic first, which seems hard to be implemented when using multinomial distribution in `numpyro`

. I really appreciate it if you can help me!