The multinomial distribution is not restricted to 2D, but it's moment assumes this is the case. This should be similar to https://github.com/pymc-devs/pymc/pull/5225