Add unit test for non-right-most batch dimension indexing
@@ -137,5 +137,57 @@ class TestMeanFieldBatchDecoupledVariationalGPBatchDim(TestBatchDecoupledVariati
return gpytorch.variational.MeanFieldVariationalDistribution
+class TestBatchDecoupledVariationalGPOtherBatchDim(TestBatchDecoupledVariationalGP, unittest.TestCase):
+ def _make_model_and_likelihood(
+ self,
+ num_inducing=16,
+ batch_shape=torch.Size([]),
+ inducing_batch_shape=torch.Size([]),
+ strategy_cls=gpytorch.variational.VariationalStrategy,
+ distribution_cls=gpytorch.variational.CholeskyVariationalDistribution,
+ constant_mean=True,
+ ):
+ class _SVGPRegressionModel(gpytorch.models.ApproximateGP):
+ def __init__(self, inducing_points):
+ variational_distribution = distribution_cls(num_inducing, batch_shape=batch_shape)
+ variational_strategy = strategy_cls(
+ self, inducing_points, variational_distribution, learn_inducing_locations=True
+ )
+ super().__init__(variational_strategy)
+ if constant_mean:
+ self.mean_module = gpytorch.means.ConstantMean(
+ batch_shape=batch_shape[:-1] + torch.Size([2]) + batch_shape[-1:]
+ )
+ self.mean_module.initialize(constant=1.0)
+ else:
+ self.mean_module = gpytorch.means.ZeroMean()
+ self.covar_module = gpytorch.kernels.ScaleKernel(
+ gpytorch.kernels.RBFKernel(batch_shape=batch_shape[:-1] + torch.Size([2]) + batch_shape[-1:]),
+ batch_shape=batch_shape[:-1] + torch.Size([2]) + batch_shape[-1:],
+ )
+
+ def forward(self, x):
+ mean_x = self.mean_module(x)
+ covar_x = self.covar_module(x)
+ latent_pred = gpytorch.distributions.MultivariateNormal(mean_x, covar_x)
+ return latent_pred
+
+ inducing_points = torch.randn(num_inducing, 2).repeat(*inducing_batch_shape, 1, 1)
+ return _SVGPRegressionModel(inducing_points), self.likelihood_cls()
+
+ @property
+ def strategy_cls(self):
+ def _batch_dim_strategy_cls(model, inducing_points, variational_distribution, learn_inducing_locations):
+ return gpytorch.variational.BatchDecoupledVariationalStrategy(
+ model, inducing_points, variational_distribution, learn_inducing_locations, mean_var_batch_dim=-2
+ )
+
+ return _batch_dim_strategy_cls
+
+ @property
+ def batch_shape(self):
+ return torch.Size([3])
+
+
if __name__ == "__main__":
unittest.main()
