Array index generated by Categorical Distribution

Questions
1

In the code below (simplification of actual model), I am attempting to select a matrix (array) from a list using indices generated by a Categorical prior distribution. I do not understand why this generates an index out of bounds error. The theano code appears to work in isolation, but not in the PyMC3 model.

from pymc3 import *
from pymc3.distributions.multivariate import Dirichlet, Multinomial
from pymc3.distributions.discrete import Categorical
import numpy as np
import warnings
import theano
import theano.tensor as tt
from theano import config, scalar


def run_MCMC3(sfs, prior, tree_matrices, tree_probabilities, draws):
    config.warn.round = False
    config.exception_verbosity='high'
    warnings.simplefilter("ignore", FutureWarning)
    tree_matrices = theano.shared(tree_matrices)


    with Model() as categorical_model:    
        tree_index = Categorical('tree_index', tree_probabilities, testval=0) #deleted testval=0        
        mx1 = tree_matrices[tree_index]
        
        probs = Dirichlet('probs', a=prior)
        conditional_probs = tt.dot(mx1, probs.T)     
        sfs_obs = Multinomial('sfs_obs', n=np.sum(sfs), p=conditional_probs, observed=sfs)
        
    with categorical_model:
        step = Metropolis()
        tune = int(draws / 5)
        trace = sample(draws, tune=tune, njobs=1, step=step) 
    
    plt.show(forestplot(trace, varnames=['probs']))
    traceplot(trace, varnames=['probs'])
    print(summary(trace))
    return trace


sfs = np.array([36, 24,  3,  9,  0])
tree_matrices = [np.array([[1., 2., 3., 4., 6.],
       [0., 0., 0., 1., 0.],
       [0., 0., 1., 0., 0.],
       [0., 1., 0., 0., 0.],
       [1., 0., 0., 0., 0.]]), np.array([[0., 2., 3., 4., 6.],
       [1., 0., 0., 1., 0.],
       [0., 0., 1., 0., 0.],
       [1., 1., 0., 0., 0.],
       [0., 0., 0., 0., 0.]]), np.array([[0., 1., 3., 4., 6.],
       [1., 1., 0., 1., 0.],
       [0., 1., 1., 0., 0.],
       [1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]]), np.array([[0., 1., 2., 4., 6.],
       [1., 1., 2., 1., 0.],
       [0., 1., 0., 0., 0.],
       [1., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0.]])]
tree_matrices = np.array(tree_matrices)
print(tree_matrices)
tree_probabilities = np.asarray([0.3, 0.2, 0.2, 0.3])
print('tree_probabilities = ', tree_probabilities)
draws = 50000
uninf_prior = np.ones(5) / 5
uninf_trace = run_MCMC3(sfs, uninf_prior, tree_matrices, tree_probabilities, draws=draws)
2

The proposal value from the Metropolis sampler of the Categorical RV tree_index went out of the support, which generate index error down stream.
The solution is not to assign Metropolis sampler to this RV, ie:

step = Metropolis([probs])

where PyMC3 will assign CategoricalGibbsMetropolis to tree_index:

Sequential sampling (2 chains in 1 job)
CompoundStep
>Metropolis: [probs]
>CategoricalGibbsMetropolis: [tree_index]
100%|██████████| 60000/60000 [00:42<00:00, 1423.08it/s]
100%|██████████| 60000/60000 [00:45<00:00, 1331.86it/s]
3

Beautiful! Thanks very much.