MutableData for boolean masks

Answer occurred to me 5 mins after posting. Will leave here for posterity.

Just breakdown the computation of mu into mu_censored and mu_uncensored and do the masking before passing to pm.MutableData:


with pm.Model():

    ...

    x_uncensored = pm.MutableData("x_uncensored", x_data[mask])
    x_censored = pm.MutableData("x_censored", x_data[~mask])
    
    mu_uncensored = x_uncensored @ beta
    mu_censored = x_censored @ beta
    ...
2 Likes