[Breaking Change] WIP: remove AdditiveStructureKernel and ProductStru…

…ctureKernel TODO: update docs / examples
cornellius-gp · Jul 2, 2024 · 9e6d35b · 9e6d35b
1 parent a049476
commit 9e6d35b
Show file tree

Hide file tree

Showing 6 changed files with 8 additions and 180 deletions.
diff --git a/docs/source/kernels.rst b/docs/source/kernels.rst
@@ -119,24 +119,12 @@ Composition/Decoration Kernels
 .. autoclass:: MultiDeviceKernel
    :members:
 
-:hidden:`AdditiveStructureKernel`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: AdditiveStructureKernel
-   :members:
-
 :hidden:`ProductKernel`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: ProductKernel
    :members:
 
-:hidden:`ProductStructureKernel`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: ProductStructureKernel
-   :members:
-
 :hidden:`ScaleKernel`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 

diff --git a/gpytorch/kernels/__init__.py b/gpytorch/kernels/__init__.py
@@ -1,6 +1,5 @@
 #!/usr/bin/env python3
 from . import keops
-from .additive_structure_kernel import AdditiveStructureKernel
 from .arc_kernel import ArcKernel
 from .constant_kernel import ConstantKernel
 from .cosine_kernel import CosineKernel
@@ -24,7 +23,6 @@
 from .piecewise_polynomial_kernel import PiecewisePolynomialKernel
 from .polynomial_kernel import PolynomialKernel
 from .polynomial_kernel_grad import PolynomialKernelGrad
-from .product_structure_kernel import ProductStructureKernel
 from .rbf_kernel import RBFKernel
 from .rbf_kernel_grad import RBFKernelGrad
 from .rbf_kernel_gradgrad import RBFKernelGradGrad
@@ -39,7 +37,6 @@
     "Kernel",
     "ArcKernel",
     "AdditiveKernel",
-    "AdditiveStructureKernel",
     "ConstantKernel",
     "CylindricalKernel",
     "MultiDeviceKernel",
@@ -61,7 +58,6 @@
     "PolynomialKernel",
     "PolynomialKernelGrad",
     "ProductKernel",
-    "ProductStructureKernel",
     "RBFKernel",
     "RFFKernel",
     "RBFKernelGrad",

diff --git a/gpytorch/kernels/additive_structure_kernel.py b/gpytorch/kernels/additive_structure_kernel.py
diff --git a/gpytorch/kernels/product_structure_kernel.py b/gpytorch/kernels/product_structure_kernel.py
diff --git a/test/examples/test_kissgp_additive_regression.py b/test/examples/test_kissgp_additive_regression.py
@@ -9,7 +9,7 @@
 
 import gpytorch
 from gpytorch.distributions import MultivariateNormal
-from gpytorch.kernels import AdditiveStructureKernel, GridInterpolationKernel, RBFKernel, ScaleKernel
+from gpytorch.kernels import GridInterpolationKernel, RBFKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ZeroMean
 
@@ -36,14 +36,12 @@ class GPRegressionModel(gpytorch.models.ExactGP):
     def __init__(self, train_x, train_y, likelihood):
         super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
         self.mean_module = ZeroMean()
-        self.base_covar_module = ScaleKernel(RBFKernel(ard_num_dims=2))
-        self.covar_module = AdditiveStructureKernel(
-            GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=1), num_dims=2
-        )
+        self.base_covar_module = ScaleKernel(RBFKernel(batch_shape=torch.Size([2])))
+        self.covar_module = GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=1)
 
     def forward(self, x):
         mean_x = self.mean_module(x)
-        covar_x = self.covar_module(x)
+        covar_x = self.covar_module(x.mT[..., None]).sum(dim=-3)
         return MultivariateNormal(mean_x, covar_x)
 
 

diff --git a/test/examples/test_kissgp_multiplicative_regression.py b/test/examples/test_kissgp_multiplicative_regression.py
@@ -10,7 +10,7 @@
 
 import gpytorch
 from gpytorch.distributions import MultivariateNormal
-from gpytorch.kernels import GridInterpolationKernel, ProductStructureKernel, RBFKernel, ScaleKernel
+from gpytorch.kernels import GridInterpolationKernel, RBFKernel, ScaleKernel
 from gpytorch.likelihoods import GaussianLikelihood
 from gpytorch.means import ConstantMean
 from gpytorch.priors import SmoothedBoxPrior
@@ -42,14 +42,12 @@ class GPRegressionModel(gpytorch.models.ExactGP):
     def __init__(self, train_x, train_y, likelihood):
         super(GPRegressionModel, self).__init__(train_x, train_y, likelihood)
         self.mean_module = ConstantMean(constant_prior=SmoothedBoxPrior(-1, 1))
-        self.base_covar_module = ScaleKernel(RBFKernel())
-        self.covar_module = ProductStructureKernel(
-            GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=1), num_dims=2
-        )
+        self.base_covar_module = ScaleKernel(RBFKernel(batch_shape=torch.Size([2])))
+        self.covar_module = GridInterpolationKernel(self.base_covar_module, grid_size=100, num_dims=1)
 
     def forward(self, x):
         mean_x = self.mean_module(x)
-        covar_x = self.covar_module(x)
+        covar_x = self.covar_module(x.mT[..., None]).prod(dim=-3)
         return MultivariateNormal(mean_x, covar_x)