Thanks for the fast reply! That’s exactly what I’m trying to do.
I can’t use p_high.argmax() here because I need to sample according to p. How do I rewrite the scan loop into a distribution? Do I need to make a new theano op, or a custom pymc3 distribution, or something else? I can’t think of a way to get around the loop because I only get the probabilities for each Categorical sample from the probabilities of the previous sample.