Improve gaussian likelihood for sampling methods with Pyro
@@ -92,11 +92,15 @@ class GaussianLikelihood(_GaussianLikelihoodBase):
y_mean = variational_dist_f.mean
if y_mean.dim() == 1:
noise = noise.squeeze(0)
- y_lazy_covar = DiagLazyTensor(var_f + noise.expand_as(var_f))
- y_dist = MultivariateNormal(y_mean, y_lazy_covar)
- if len(y_dist.batch_shape):
- y_dist = y_dist.__class__(
- y_dist.mean.contiguous().view(-1), BlockDiagLazyTensor(y_dist.lazy_covariance_matrix)
+
+ y_dist = pyro.distributions.Independent(
+ pyro.distributions.Normal(y_mean, (var_f + noise.expand_as(var_f)).sqrt()),
+ reinterpreted_batch_ndims=y_mean.dim(),
)
- y_obs = y_obs.view_as(y_dist.mean)
+
+ # See if we're using a sampled GP distribution
+ # Samples will occur in the first batch dimension
+ sample_shape = y_dist.shape()[:-y_obs.dim()]
+ y_obs = y_obs.expand(y_dist.shape())
+ with pyro.poutine.scale(scale=float(1. / sample_shape.numel())):
pyro.sample(name_prefix + "._training_labels", y_dist, obs=y_obs)
