In the simplest case of a categorical you hape a vector of probabilities for a scalar output categorical(p=[0.5 0.5]) which can yield a single 0 or 1. Instead categorical(p=[0.3, 0.3, 0.4]) can yield a single 0, 1, or 2.
The last axis of p gives you the event probabilities. If you have a (3, 2) p matrix you get 3 scalars in {0, 1}, if you have a (2, 3) p matrix you get 2 scalars in {0, 1, 2}. Does that help?