argchecks in spectral mixture kernel
@@ -67,20 +67,14 @@ class SpectralMixtureKernel(Kernel):
.. _Gaussian Process Kernels for Pattern Discovery and Extrapolation:
https://arxiv.org/pdf/1302.4245.pdf
"""
- # TODO: add equation to docs
-
def __init__(
self,
num_mixtures=None,
ard_num_dims=1,
- batch_size=1,
- active_dims=None,
- eps=1e-6,
+ batch_shape=torch.Size([1]),
mixture_scales_prior=None,
mixture_means_prior=None,
mixture_weights_prior=None,
- param_transform=softplus,
- inv_param_transform=None,
**kwargs
):
if num_mixtures is None:
@@ -88,19 +82,17 @@ class SpectralMixtureKernel(Kernel):
if mixture_means_prior is not None or mixture_scales_prior is not None or mixture_weights_prior is not None:
logger.warning("Priors not implemented for SpectralMixtureKernel")
+ if len(batch_shape) > 1:
+ raise RuntimeError("SpectralMixtureKernel does not yet support multiple batch dimensions.")
+
# This kernel does not use the default lengthscale
- super(SpectralMixtureKernel, self).__init__(
- active_dims=active_dims, param_transform=param_transform, inv_param_transform=inv_param_transform
- )
+ super(SpectralMixtureKernel, self).__init__(ard_num_dims=ard_num_dims, batch_shape=batch_shape, **kwargs)
self.num_mixtures = num_mixtures
- self.batch_size = batch_size
- self.ard_num_dims = ard_num_dims
- self.eps = eps
self.register_parameter(
- name="raw_mixture_weights", parameter=torch.nn.Parameter(torch.zeros(self.batch_size, self.num_mixtures))
+ name="raw_mixture_weights", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, self.num_mixtures))
)
- ms_shape = torch.Size([self.batch_size, self.num_mixtures, 1, self.ard_num_dims])
+ ms_shape = torch.Size([*self.batch_shape, self.num_mixtures, 1, self.ard_num_dims])
self.register_parameter(name="raw_mixture_means", parameter=torch.nn.Parameter(torch.zeros(ms_shape)))
self.register_parameter(name="raw_mixture_scales", parameter=torch.nn.Parameter(torch.zeros(ms_shape)))
@@ -206,18 +198,21 @@ class SpectralMixtureKernel(Kernel):
return x1_.unsqueeze(-2), x2_.unsqueeze(-3)
def forward(self, x1, x2, **params):
- batch_size, n, num_dims = x1.size()
- _, m, _ = x2.size()
+ if x1.dim() > 3 or x2.dim() > 3:
+ raise RuntimeError("SpectralMixtureKernel does not yet support multiple batch dimensions.")
+
+ batch_shape = x1.shape[:-2]
+ n, num_dims = x1.shape[-2:]
if not num_dims == self.ard_num_dims:
raise RuntimeError(
"The SpectralMixtureKernel expected the input to have {} dimensionality "
"(based on the ard_num_dims argument). Got {}.".format(self.ard_num_dims, num_dims)
)
- if not batch_size == self.batch_size:
+ if not batch_shape == self.batch_shape:
raise RuntimeError(
"The SpectralMixtureKernel expected the input to have a batch_size of {} "
- "(based on the batch_size argument). Got {}.".format(self.batch_size, batch_size)
+ "(based on the batch_size argument). Got {}.".format(self.batch_shape, batch_shape)
)
# Expand x1 and x2 to account for the number of mixtures
