MutableData for boolean masks

I’m setting up a model involving censored observations. I’d like to use pm.MutableData for out of sample predictions. The basic structure of the model is:

normal_lcdf = pm.distributions.dist_math.normal_lcdf

with pm.Model():
    x = pm.MutableData("x", x_data)  # (n_obs, n_features) predictors
    y = pm.MutableData("y", y_data)  # (n_obs,) responses
    beta = ...  # (n_features,) coefficients
    sigma = ...  # std dev
    mu = x @ beta  # (n_obs,) modelled responses

    # Uncensored observations
    mask = np.array([False, True, ...])  # (n_obs,) boolean array
    uncensored = pm.MutableData("uncensored", mask)
        "obs_uncensored", mu=mu[uncensored], sigma=sigma, observed=y[uncensored]

    # Censored observations
    censored = pm.MutableData("censored", ~mask)
            mu=mu[censored], sigma=sigma, x=y[censored]

This throws: TypeError: index must be integers or a boolean mask presumably because you
can’t index with pm.MutableData.

I thought about masking before passing to pm.MutableData, like:

mask = np.array([False, True, ...])  # (n_obs,) boolean array

with pm.Model():

    y_uncensored = pm.MutableData("y_uncensored", y_data[mask])
    y_censored = pm.MutableData("y_censored", y_data[~mask])


But I would still have to index into mu, which would mean redefining mu for new data.

Is there a way to implement masks using pm.MutableData? Is there another way to go
about this?

Thanks in advance for any suggestions

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