How to implement Zero-Inflated Multinomial?

I just ran this version and… it actually takes more time to sample :sweat_smile:
Did I make a mistake somewhere? Or it’s just that this model can’t be optimized anymore?