Simplify prior mode vs posterior mode checks
@@ -75,9 +75,18 @@ class ExactGP(GP):
res = super(ExactGP, self).__call__(*inputs, **kwargs)
return res
+ # Prior mode
+ elif self.train_inputs is None or self.train_targets is None:
+ full_inputs = args
+ full_output = super(ExactGP, self).__call__(*full_inputs, **kwargs)
+ if settings.debug().on():
+ if not isinstance(full_output, MultivariateNormal):
+ raise RuntimeError("ExactGP.forward must return a MultivariateNormal")
+ full_mean, full_covar = full_output.mean, full_output.lazy_covariance_matrix
+ return full_output.__class__(full_mean, full_covar)
# Posterior mode
else:
- if settings.debug.on() and self.train_inputs is not None:
+ if settings.debug.on():
if all(torch.equal(train_input, input) for train_input, input in zip(train_inputs, inputs)):
warnings.warn(
"The input matches the stored training data. Did you forget to call model.train()?", UserWarning
@@ -85,9 +94,7 @@ class ExactGP(GP):
# Exact inference
non_batch_train = False
- if self.train_inputs is None:
- full_inputs = args
- else:
+
if all(tin.dim() == 2 for tin in self.train_inputs):
# Train inputs were non-batch
non_batch_train = True
@@ -113,7 +120,6 @@ class ExactGP(GP):
num_train = 0
train_targets = None
- if self.train_targets is not None:
train_targets = self.train_targets
# If we expanded the train_inputs, we need to do the same for the train_targets
@@ -172,8 +178,5 @@ class ExactGP(GP):
else:
# Standard multitask
predictive_mean = predictive_mean.view(-1, num_tasks).contiguous()
- else:
- predictive_mean = full_mean
- predictive_covar = full_covar
return full_output.__class__(predictive_mean, predictive_covar)
