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
...