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!