The code you posted appears to generate samples from a markov chain with a known transition matrix, then computes the total share of state 2 observed. You can accomplish this by making a pmx.DiscreteMarkovChain RV with P=transition_matrix, calling pm.sample_prior_predictive, then computing the share of state 2 in the samples.