TypeError in NUTS

That seems to work.

Right before you posted that, I figured out that this error goes away if I do

with pm.Model() as model:
    ...
    # use categorical because Metropolis gives invalid proposal value for DiscreteUniform sometimes
    probs = pm.DiscreteUniform('probs', lower=0, upper=499, shape=(n_train - n_obs))
    p = tt.extra_ops.to_one_hot(probs.reshape((n_train - n_obs, 1)), 500)
    ...
    # sample using the auto assigned sampler
    trace_discrete = pm.sample(500, tune=1000)

And it will start running, but somewhere around step 200 it throws a very long and unintelligible IndexError.