Simple Exponential Survival Funciton

This line, it looks like forgot to index the relevant log_rate

Should be

y_obs  = pm.Exponential('y_obs',lam=log_rate[adopted==0], observed=days_to_event[adopted==0])

Unrelated to the shape issues, you could also use pm.Censored https://www.pymc.io/projects/docs/en/stable/api/distributions/generated/pymc.Censored.html

with pm.Model() as cat_censored_m:
    days_to_event = np.array(cats['days_to_event'])
    adopted = np.array([0 if x == 'Adoption' else 1 for x in cats['out_event']])
    black = np.array([1 if x == 'Black' else 0 for x in cats['color']])
    
    a = pm.Normal('a', 0, 1,shape=2)
    log_rate = a[black]

    upper_censoring = days_to_event.copy()
    upper_censoring[adopted==0] = np.inf  

    y_obs = pm.Censored(
        'y_obs', 
        pm.Exponential.dist(lam=log_rate),
        upper=upper_censoring, 
        observed=days_to_event,
    )
   
    cat_trace = pm.sample() 
1 Like