Working code fails with PyMC 5.16

1

Hi,

I have some code that I have been running with no issues for a few months.
I created a new environment on a new machine that uses PyMC 5.16.2 and Pytensor 2.25.1.
The code works fine when using PyMC 5.14.0 and Pytensor 2.20.0.
I followed all the recommended steps to install pymc, I tested this with both macos and linux, I tested both python 3.11 and 3.12, I also tested both the pymc and numpyro samplers.

I have attached a colab [notebook] that reproduces the error (Google Colab).

The code creates a reinforcement learning model and estimates its parameter for multiple subjects at once - the code was adapted from this pymc tutorial.

The error occurs in the following function, which is passed to a pytensor scan operation to compute the logp of a CustomDist.

def update_belief(choices, outcomes, belief, choice_probs,
                  n_subj,
                  alpha, beta):
    c = (choices)
    nc = (1 - choices)

    all_subjs = pt.arange(n_subj)

    #### ERROR here
    #### "TypeError: Trying to increment a 1-dimensional subtensor with a 3-dimensional value."
    # compute choice probabilities
    choice_probs = pt.set_subtensor(choice_probs[all_subjs],
                                    pt.exp(beta[all_subjs] * belief[all_subjs,c])
                                      / (pt.exp(beta[all_subjs] * belief[all_subjs,0]) + pt.exp(beta[all_subjs] * belief[all_subjs,1])))

    # update beliefs of chosen arm
    belief = pt.set_subtensor(belief[all_subjs,c],
                              belief[all_subjs,c] + alpha[all_subjs] * (outcomes[all_subjs,c] - belief[all_subjs,c]))

    # update beliefs of not chosen arm
    belief = pt.set_subtensor(belief[all_subjs,nc],
                              belief[all_subjs,nc])

    return belief, choice_probs

def rl_logp(observed, alpha, beta):
    # unpack variables
    choices = observed[:,:,0].astype('int32')
    outcomes = observed[:,:,1:].astype('int32')
    n_subj = choices.shape[1]

    choices_ = pt.as_tensor_variable(choices, dtype='int32')
    outcomes_ = pt.as_tensor_variable(outcomes, dtype='int32')

    beliefs = pt.zeros((n_subj,2), dtype='float64')   # [n_subj x 2]
    choice_probs_ = pt.zeros(n_subj, dtype='float64') # [n_subj x 1]

    [beliefs_pymc, choice_probs_pymc], updates = scan(
        fn=update_belief,
        sequences=[choices_, outcomes_],
        non_sequences=[n_subj, alpha, beta],
        outputs_info=[beliefs, choice_probs_],
        profile=True
    )

    # pymc expects the standard (positive) log-likelihood
    ll = pt.sum(pt.log(choice_probs_pymc), axis=0)
    return ll

with pm.Model() as model:
    ''' Fix inverse temperature to 4 and only fit alpha in this example '''
    observed = np.concatenate([simulated_choices[:,:,None], outcomes], axis=2)

    # hyper priors for alpha
    alpha_a = pm.Normal('alpha_a', mu=0, sigma=1)
    alpha_b = pm.HalfNormal('alpha_b', sigma=0.2)

    # non-centered parametrisation
    # priors
    alpha_raw = pm.Normal('alpha_raw', mu=0, sigma=1, shape=n_subj)
    alpha = pm.Deterministic('alpha', pm.math.invprobit(alpha_a + alpha_b * alpha_raw))

    # fix beta to 4
    beta = pm.Deterministic('beta', 4 * pm.math.ones_like(np.ones(n_subj)))

    logp = pm.CustomDist('logp', alpha, beta,
                         logp=rl_logp, observed=observed)

What is the correct way to index through the choice probabilities and assign them to the choice_probs array?
I tried to have a look at the github releases but I cannot pinpoint this issue to any specific commit, can any one from the development team help with this?

Thanks!

2

The error says that in the following line

choice_probs = pt.set_subtensor(
    choice_probs[all_subjs],
    pt.exp(beta[all_subjs] * belief[all_subjs,c])
     / (pt.exp(beta[all_subjs] * belief[all_subjs,0]) + pt.exp(beta[all_subjs] * belief[all_subjs,1]))
)

choice_probs[all_subjs] is a vector
And the update is a tensor3. Which one is correct?

3

Hi Ricardo,

I am not sure which one is correct in this case - I imagine that some automatic broadcasting/conversion handled this up to v5.14.0 and I never worried about it since the results were matching the numpy implementation.

I am trying to replicate the following code, with scan “looping” through the trials and computing the choice probabilities at each trial.
But, I am not really sure on how to change the pytensor code to match this.

n_subj = 10
n_trials = 20

beta = np.ones((n_subj))
choice_probs = np.zeros((n_trials, n_subj))
belief = np.random.rand(n_trials, n_subj, 2)

for t in range(n_trials):
    choice_probs = np.exp(beta[:] * belief[t,:,0]) / (np.exp(beta[:] * belief[t,:,0]) + np.exp(beta[:] * belief[t,:,1]))