Unexpected nan in Normal and HalfNormal variables (debugging with printout)

Encountered an unexpected nan in a hierarchical model. These are the lines related to the nan error; I list the full code after the error printout.

alpha = pm.Normal("alpha", mu=np.array([np.log(0.002), np.log(0.002), np.log(0.5/(1-0.5)), sigma=1, shape=3)
print_alpha = Print("alpha")(alpha)
omega = pm.HalfNormal("omega", sigma=1, shape=3)
print_omega = Print("omega")(omega)

from utilities import *
import arviz as az
import pymc as pm
# Function argument shapes: 
# patient dictionary contains N_patients patients

def individual_model(patient_dictionary, name):
    df = pd.DataFrame(columns=["patient_id", "mprotein_value", "time"])
    for ii in range(len(patient_dictionary)):
        patient = patient_dictionary[ii]
        mprot = patient.Mprotein_values
        times = patient.measurement_times
        for jj in range(len(mprot)):
            single_entry = pd.DataFrame({"patient_id":[ii], "mprotein_value":[mprot[jj]], "time":[times[jj]]})
            df = pd.concat([df, single_entry], ignore_index=True)
    group_id = df["patient_id"].tolist()
    Y_flat_no_nans = np.array(df["mprotein_value"].tolist())
    assert min(Y_flat_no_nans) >= 0
    t_flat_no_nans = np.array(df["time"].tolist())
    assert min(t_flat_no_nans) >= 0

    N_patients = len(patient_dictionary)
    yi0 = np.zeros(N_patients)
    for ii in range(N_patients):
        yi0[ii] = max(patient_dictionary[ii].Mprotein_values[0], 1e-5)
    assert min(yi0) > 0 #Strictly greater than zero required because we log transform it for the log prior of psi 

    with pm.Model() as individual_model:
        # Observation noise (std)
        sigma_obs = pm.HalfNormal("sigma_obs", sigma=1)
        print_sigma_obs = Print("sigma_obs ")(sigma_obs)

        # alpha
        alpha = pm.Normal("alpha",  mu=np.array([np.log(0.002), np.log(0.002), np.log(0.5/(1-0.5))]),  sigma=1, shape=3)
        print_alpha = Print("alpha")(alpha)

        # Latent variables theta
        omega = pm.HalfNormal("omega",  sigma=1, shape=3) # Patient variability in theta (std)
        print_omega = Print("omega")(omega)

        theta_rho_s = pm.Normal("theta_rho_s", mu= print_alpha[0], sigma=print_omega[0], shape=N_patients)
        theta_rho_r = pm.Normal("theta_rho_r", mu= print_alpha[1], sigma=print_omega[1], shape=N_patients)
        theta_pi_r  = pm.Normal("theta_pi_r",  mu= print_alpha[2], sigma=print_omega[2], shape=N_patients)

        # psi: True M protein at time 0
        xi = pm.HalfNormal("xi", sigma=1)
        log_psi = pm.Normal("log_psi", mu=np.log(yi0), sigma=xi, shape=N_patients)
        psi = pm.Deterministic("psi", np.exp(log_psi))

        # Transformed latent variables 
        rho_s = pm.Deterministic("rho_s", -np.exp(theta_rho_s))
        rho_r = pm.Deterministic("rho_r", np.exp(theta_rho_r))
        pi_r  = pm.Deterministic("pi_r", 1/(1+np.exp(-theta_pi_r)))

        mu_Y = psi[group_id] * (pi_r[group_id]*np.exp(rho_r[group_id]*t_flat_no_nans) + (1-pi_r[group_id])*np.exp(rho_s[group_id]*t_flat_no_nans))

        # Likelihood (sampling distribution) of observations
        Y_obs = pm.Normal("Y_obs", mu=mu_Y, sigma=print_sigma_obs, observed=Y_flat_no_nans)
    return individual_model