Model finishes sampling then raises dimension error

Hi there. I’ve encountered an error that seems very strange to me. When I sample the model below, it samples nicely (using the pymc sampler), and gets to the full 2000 draws on each chain, then throws the error ValueError: conflicting sizes for dimension 'participant_id': length 1 on the data but length 40 on coordinate 'participant_id' rather than returning the inference data. I can’t fathom what’s going on, since the model compiles and samples fine before throwing this error. Can someone point me in the right direction please?

def belief_update(llr, is_explicit_switch, explicit_states, belief_tm1, hazard_rate, log_odds_exp_switch):
    """
    The equation used here comes from Normative evidence accumulation in unpredictable environments by Glaze et al, 2015. Data should be coded such that centre a is negative and centre b is positive.

    Args:
        belief_tm1 (float): The belief state after the last update
        llr (float): The log likelihood ratio of the evidence for centre b over centre a
        hazard_rate (float): The probability of a switch each time

    Returns:
        float: The belief state (odds of centre b) given the new evidence.
    """
    log_prior = (
        belief_tm1 + 
        pt.log((1 - hazard_rate) / hazard_rate + pt.exp(-belief_tm1)) - 
        pt.log((1 - hazard_rate) / hazard_rate + pt.exp(belief_tm1))
    )

    new_belief = (log_prior + llr) * (1 - is_explicit_switch) + (explicit_states * log_odds_exp_switch) * is_explicit_switch
    # The multiplication by is_explicit_switch is an easier-sampling alternative to a switch statement

    return new_belief

def perceptual_model(initial_belief, log_likelihood_ratios, is_explicit_switch, explicit_states, hazard_rate, log_odds_exp_switch):
    belief_states, _ = pytensor.scan(
        fn=belief_update,
        sequences=[log_likelihood_ratios, is_explicit_switch, explicit_states],
        outputs_info=[{'initial': [initial_belief]}],
        non_sequences=[hazard_rate, log_odds_exp_switch],
    )

    return belief_states[:, :, 0]


coords = {
    'trial_number': np.arange(MAX_TRIALS),
    'participant_id': participant_ids,
}
with pm.Model(coords=coords) as observer_model:
    logp_centre_a = pt.log(stats.norm.pdf(target_coins, loc=-0.05, scale=0.05))
    logp_centre_b = pt.log(stats.norm.pdf(target_coins, loc=0.05, scale=0.05))
    log_likelihood_ratios = logp_centre_b - logp_centre_a

    initial_belief = pm.Data(
        'initial_belief',
        initial_belief,
        dims='participant_id'
    )

    group_hazard_rate = pm.Beta('group_hazard_rate', alpha=2, beta=2)
    logit_hazard_rate_sigma = pm.HalfNormal('hazard_rate_sigma', sigma=2)
    logit_hazard_rate_offset = pm.Normal(
        'logit_hazard_rate_offset', mu=0, sigma=1, dims='participant_id'
    )
    hazard_rate = pm.Deterministic(
        'hazard_rate',
        pm.math.sigmoid(
            pm.math.logit(group_hazard_rate) + logit_hazard_rate_offset * logit_hazard_rate_sigma
        ),
        dims='participant_id'
    )

    group_log_odds_exp_switch = pm.Normal('group_log_odds_exp_switch', mu=1, sigma=2)
    sigma_log_odds_exp_switch = pm.HalfNormal('sigma_log_odds_exp_switch', sigma=2)
    log_odds_exp_switch_offset = pm.Normal(
        'log_odds_exp_switch_offset', mu=0, sigma=1, dims='participant_id'
    )
    log_odds_exp_switch = pm.Deterministic(
        'log_odds_exp_switch',
        group_log_odds_exp_switch + log_odds_exp_switch_offset * sigma_log_odds_exp_switch,
        dims='participant_id'
    )
    
    generative_centres = pm.Data(
        'generative_centres',
        generative_centres,
        dims=('trial_number', 'participant_id'),
        )
    is_explicit_switch = pm.Data(
        'is_explicit_switch',
        is_explicit_switch,
        dims=('trial_number', 'participant_id'),
        )

    beliefs = pm.Deterministic(
        'beliefs',
        perceptual_model(initial_belief, log_likelihood_ratios, is_explicit_switch, generative_centres, hazard_rate, log_odds_exp_switch),
        dims=('trial_number', 'participant_id')
    )

    #----Determine participant belief----
    centre_separation = 0.1
    x_a = -centre_separation/2
    x_b = centre_separation/2

    stim_coins = pm.Data(
        'stim_coins', stim_coins, dims=('trial_number', 'participant_id')
        )
    actual_responses = pm.Data(
        'actual_responses', actual_responses, dims=('trial_number', 'participant_id')
        )

    gain=GAIN

    guess_if_a = x_a + (stim_coins - x_a) * gain
    guess_if_b = x_b + (stim_coins - x_b) * gain
    
    inferred_centre = pm.Deterministic(
        'inferred_belief',
        pt.where(
            pt.abs(guess_if_a - actual_responses) < pt.abs(guess_if_b - actual_responses),
            0, 1
            ),
        dims=('trial_number', 'participant_id')
    )

    group_p_random_error = pm.Beta('group_p_random_error', mu=0.1, nu=5)
    sigma_logit_p_random_error = pm.HalfNormal('sigma_logit_p_random_error', sigma=2)
    p_random_error_offset = pm.Normal(
        'p_random_error_offset', mu=0, sigma=1, dims='participant_id'
    )
    p_random_error = pm.Deterministic(
        'p_random_error',
        pm.math.sigmoid(
            pm.math.logit(group_p_random_error) + p_random_error_offset * sigma_logit_p_random_error
        ),
        dims='participant_id'
    )

    belief_as_p = pt.sigmoid(beliefs) # Convert log odds to probability

    p_respond_b = (
        (1 - p_random_error) * belief_as_p + p_random_error * 0.5
    ) # Belief weighted response + random error (50/50 guess)

    observed_mask = pt.as_tensor(observed_mask, dims=('trial_number', 'participant_id'))

    def logp_likelihood_fn(response, belief, is_observed):
        return pt.switch(
            pt.eq(is_observed, 1),
            pm.logp(pm.Bernoulli.dist(p=belief), response),
            pm.logp(pm.DiracDelta.dist(-1), -1)
        )

    likelihood_dist = pm.CustomDist(
        'likelihood',
        p_respond_b,
        observed_mask,
        logp=logp_likelihood_fn,
        observed=inferred_centre,
        dims=('trial_number', 'participant_id')
    )

Hey @ddg.palmer, I believe what is happening is after your model is done sampling it is converting your trace into an InferenceData object. However, it seems that there is a mismatch somewhere between the shape you specified with the dims argument and the actual shape of the object. I can’t tell my looking at it where this is happening without the shapes of your data. If you tell me the shapes I would be happy to help you debug it.

1 Like

Thanks for your reply, @Dekermanjian. You pointed me right to the problem.

For anyone for whom this might be useful in the future, the issue was that the return of the scan function, beliefs, had shape (600, 1), when I expected it to have shape (600, 40). The root cause of that turned out to me having indirectly wrapped the initial_belief variable in brackets twice in the process of passing it to scan. I changed the outputs_info argument of scan from outputs_info=[{'initial': [initial_belief]}] to outputs_info=[initial_belief] and removed the slice from the return function, and the model now samples properly and gives the expected result for my Bayesian observer model.

Thanks again for your help.

David

2 Likes

I am glad you were able to figure it out!