Change index kernel in mixed factory to just be multitask (#799)

Summary:
Pull Request resolved: #799

The index kernel for the mixed factory now acts more like a multitask kernel as opposed to the cateogrical kernel which acts more like a true categorical kernel.

Reviewed By: tymmsc

Differential Revision: D74489044

fbshipit-source-id: 1c03c8bb5fd2065da3ed7f3d02f73654666bbb39

Author: JasonKChow

aepsych/factory/mixed.py

@@ -109,13 +109,10 @@ def _make_covar_module(self) -> gpytorch.kernels.Kernel:
                         ),
                     )
                 )
-            add_kernel = gpytorch.kernels.AdditiveKernel(
-                deepcopy(cont_kernel), *deepcopy(discrete_kernels)
-            )
-            prod_kernel = gpytorch.kernels.ProductKernel(
+
+            return gpytorch.kernels.ProductKernel(
                 deepcopy(cont_kernel), *deepcopy(discrete_kernels)
             )
-            return add_kernel * prod_kernel
         elif self.discrete_kernel == "categorical":
             constraint = gpytorch.constraints.GreaterThan(lower_bound=1e-4)
             discrete_kernel = botorch.models.kernels.CategoricalKernel(

tests/test_mean_covar_factories.py

@@ -703,32 +703,16 @@ def test_mixed_from_config(self):
         self.assertEqual(model.dim, 4)
         self.assertIsInstance(covar, gpytorch.kernels.ProductKernel)
 
-        # Check the additive part
-        add_kernel = covar.kernels[0]
-        self.assertIsInstance(add_kernel.kernels[0], gpytorch.kernels.RBFKernel)
-        self.assertSequenceEqual(add_kernel.kernels[0].active_dims, (0, 3))
-        self.assertEqual(len(add_kernel.kernels[1:]), 2)
-        for kernel, index, rank in zip(add_kernel.kernels[1:], (1, 2), (2, 3)):
-            self.assertIsInstance(kernel, gpytorch.kernels.IndexKernel)
-            self.assertEqual(kernel.active_dims.item(), index)
-            self.assertEqual(kernel.covar_factor.shape[1], rank)
-
-        # Check the product part
-        cont_kernel = covar.kernels[1]
+        cont_kernel = covar.kernels[0]
         self.assertIsInstance(cont_kernel, gpytorch.kernels.RBFKernel)
         self.assertSequenceEqual(cont_kernel.active_dims, (0, 3))
 
-        index_kernels = covar.kernels[2:]
+        index_kernels = covar.kernels[1:]
         for kernel, index, rank in zip(index_kernels, (1, 2), (2, 3)):
             self.assertIsInstance(kernel, gpytorch.kernels.IndexKernel)
             self.assertEqual(kernel.active_dims.item(), index)
             self.assertEqual(kernel.covar_factor.shape[1], rank)
 
-        # Check there's copies and not duplicates
-        self.assertNotEqual(add_kernel.kernels[0], cont_kernel)
-        self.assertNotEqual(add_kernel.kernels[1], index_kernels[0])
-        self.assertNotEqual(add_kernel.kernels[2], index_kernels[1])
-
     def test_mixed_acquisition(self):
         def f_1d(x):
             """