Learning a Categorical Variable with TensorFlow Probability

Question

I would like to use TFP to write a neural network where the output are the probabilities of a categorical variable with 3 classes, and train it using the negative log-likeli