Questions about hierarchical time-series modeling with missing data

Hi! I am new to Pyro and I have some questions modeling a hierarchical time-series model.

My dataset consists of several time series data sampled at different locations. However, data for some of the locations are completely missing. I would like to impute the missing data with a hierarchical Bayesian model. In particular, I would like to start with a linear model with one covariate.

I am basically following the examples in the Pyro documentations on forecasting. However, the examples only consider data at one location, while I need to model multiple locations together. In addition, the examples focus on predicting for future times, while I need to “predict” for new locations. Hence, the examples do not apply directly.

Following is my attempt:

def model(self, zero_data, covariates):
        data_size, duration, feature_dim = covariates.shape

        bias_scale = pyro.sample("bias_scale", dist.LogNormal(-5,5))
        weight_scale = pyro.sample("weight_scale", dist.LogNormal(-5,5))

        with pyro.plate("load", data_size, dim=-2):
            bias = pyro.sample("bias", dist.Normal(0, bias_scale))
            weight = pyro.sample("weight", dist.Normal(0, weight_scale).expand([feature_dim]).to_event(1))
        prediction = bias + (weight * covariates).sum(-1, keepdim=True)

        noise_scale = pyro.sample("noise_scale", dist.LogNormal(-5, 5).expand([1]).to_event(1))
        noise_dist = dist.Normal(0, noise_scale)
        
        self.predict(noise_dist, prediction)

The issue with this attempt is the “prediction”. Since my “covariates” has the shape (locations, duration, features), it does not work correctly with the “*” operator. In addition, I don’t know if my way of modeling the noise would result in independent noise wrt different locations and different time or not.

I am sorry if my expression is not that clear, and please tell me if I need to provide more details of this problem. I appreciate any suggestions and advice.

Thank you very much!