This is usually the source of problem, as a high dimensional Z makes the sampling difficult, especially mixing NUTS with other samplers.
The main strategy for these kinds of problem is to marginalize out the discrete variable Z, by expressing the likelihood as a mixture distribution.
Another approach I sometimes take is to express the indexing matrix Z as a weight matrix between 0 and 1 (and summed to 1). It should not be confused with the marginalization although the result is quite similar and it is often a bit easier to rewrite your model. See eg this post: Indexing using a free random variable - #16 by junpenglao