Problem of fitting expected utility model in risky choice

Questions v5
1

Dear pymc experts,

I am trying to use pymc to fit a simple expected utility model in risky choice task.
expect utility function is simple: u = o**alpha * p, where o is the objective reward/money, p is the probability and alpha is the parameter.
I confront a very weird/puzzling problem that after the tune/burn-in stage, every sample is divergent sample and the R-hat is extraordinary high. However, if i delete the alpha parameter now the model becomes a expected value model, the sample is perfectly fine. It seems that adding the power function brings a serious problem… Here is my code:

with pm.Model(coords=coords) as m1:
     
    alpha_mu = pm.Normal('alpha_mu',0,1)
    gamma_mu = pm.Normal('gamma_mu',0,1)


    
    alpha_s = pm.HalfNormal('alpha_s',3)
    gamma_s = pm.HalfNormal('gamma_s',3)
    
    ## indiviual parameter
    alpha_i_raw = pm.Normal('alpha_i',mu = 0, sigma = 1, dims = "subj")
    gamma_i_raw = pm.Normal('gamma_i',mu = 0, sigma = 1, dims = "subj")
    
    ## matt trick
    alpha = 2*pm.math.invlogit(alpha_mu + alpha_s * alpha_i_raw)  ## make sure the range of parameter locates in [0,2]
    gamma = 5*pm.math.invlogit(gamma_mu + gamma_s * gamma_i_raw) 

    ## compute utility
    u_r = p_r * (o_r**alpha[rep_count])  - o_c**alpha[rep_count] 
    llf = pm.math.invlogit(u_r * gamma[rep_count])
    y = pm.Bernoulli('llf',p=llf,observed=choice_)
    eu_m1_r = pm.sample(draws=4000,tune=2000,chains=4,cores=4)
2

Can you provide a runnable example? It will be easier to help if someone can run it and see the problem you are seeing.