Just as a heads up, this is not ideal, since you don’t actually want to be explicitly sampling from the Multinomial distribution. It better to only be evaluating the log-likelihood.
The way to do this is to do what you were doing before.
Also note that passing in raw probabilities will incur underflow, you’ll want to pass in logits instead.
I posted a simplified example of this without the normal regression here