def forward(self, X, Y=None):
# Propagate samples through the layers.
F = X
for i, gp in enumerate(self.gps):
if isinstance(F, MultivariateNormal):
# Sample from independent Gaussian
m = F.mean.reshape(-1, gp.input_dims)
s = F.stddev.reshape(-1, gp.input_dims)
F = Normal(loc=m, scale=s).rsample()
if i > 0 and self.add_input:
# Add input to layers after the first one
assert F.shape == X.shape
F.add_(X)
# Get posterior distribution
F = gp(F, Y=Y, likelihood=self.likelihood)
if isinstance(self.likelihood, MultitaskGaussianLikelihood):
D = self.likelihood.num_tasks
else:
D = 1
# If our likelihood is one-dimensional but we are outputting multiple
# dimensions, we sum the dimensions together
if isinstance(F, MultitaskMultivariateNormal) and D == 1:
mean = F.mean.sum(axis=-1)
variance = F.variance.sum(axis=-1)
F = MultivariateNormal(mean, DiagLazyTensor(variance))
return F
