Simple markov chain with aesara scan

Questions v5
11

I think the error you see may be coming from the sampler proposing values out of bounds for the bernoulli variable, like p=transition_probs[2]. Again this is only an issue if you are not conditioning on observed data, but actually sampling the scan sequence.

If you specify a better sampler manually, it seems to work:

import numpy as np
import aesara
import pymc as pm

k = 10

with pm.Model() as markov_chain:
    
    transition_probs = pm.Uniform('transition_probs', lower=0, upper=1, shape = 2)
    initial_state = pm.Bernoulli('initial_state', p = 0.5)
    
    def transition(previous_state, transition_probs, old_rng):
        p = transition_probs[previous_state]
        next_rng, next_state = pm.Bernoulli.dist(p = p, rng=old_rng).owner.outputs
        return next_state, {old_rng: next_rng}

    rng = aesara.shared(np.random.default_rng())
    mc_chain, updates = aesara.scan(fn=transition,
                                  outputs_info=dict(initial = initial_state),
                                  non_sequences=[transition_probs, rng],
                                  n_steps=k)
    assert updates
    markov_chain.register_rv(mc_chain, name="mc_chain", initval="prior")

with markov_chain:
    pm.sample(chains=1, step=pm.BinaryMetropolis([mc_chain]))