More explicit prior shape broadcasting

gpytorch/kernels/arc_kernel.py

@@ -121,25 +121,24 @@ def __init__(
         self.register_parameter(
             name="raw_angle", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, self.last_dim)),
         )
+        self.register_constraint("raw_angle", angle_constraint)
+
         if angle_prior is not None:
             self.register_prior(
                 "angle_prior", angle_prior, lambda m: m.angle, lambda m, v: m._set_angle(v),
             )
 
-        self.register_constraint("raw_angle", angle_constraint)
-
         self.register_parameter(
             name="raw_radius", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, self.last_dim)),
         )
+        radius_constraint = Positive()
+        self.register_constraint("raw_radius", radius_constraint)
 
         if radius_prior is not None:
             self.register_prior(
                 "radius_prior", radius_prior, lambda m: m.radius, lambda m, v: m._set_radius(v),
             )
 
-        radius_constraint = Positive()
-        self.register_constraint("raw_radius", radius_constraint)
-
         self.base_kernel = base_kernel
         if self.base_kernel.has_lengthscale:
             self.base_kernel.lengthscale = 1

gpytorch/kernels/cosine_kernel.py

@@ -65,6 +65,7 @@ def __init__(self, period_length_prior=None, period_length_constraint=None, **kw
 
         if period_length_constraint is None:
             period_length_constraint = Positive()
+        self.register_constraint("raw_period_length", period_length_constraint)
 
         if period_length_prior is not None:
             self.register_prior(
@@ -74,8 +75,6 @@ def __init__(self, period_length_prior=None, period_length_constraint=None, **kw
                 lambda m, v: m._set_period_length(v),
             )
 
-        self.register_constraint("raw_period_length", period_length_constraint)
-
     @property
     def period_length(self):
         return self.raw_period_length_constraint.transform(self.raw_period_length)

gpytorch/kernels/index_kernel.py

@@ -53,11 +53,11 @@ def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs)
             name="covar_factor", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks, rank))
         )
         self.register_parameter(name="raw_var", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))
+        self.register_constraint("raw_var", var_constraint)
+
         if prior is not None:
             self.register_prior("IndexKernelPrior", prior, lambda m: m._eval_covar_matrix())
 
-        self.register_constraint("raw_var", var_constraint)
-
     @property
     def var(self):
         return self.raw_var_constraint.transform(self.raw_var)

gpytorch/kernels/kernel.py

@@ -166,13 +166,13 @@ def __init__(
                 name="raw_lengthscale",
                 parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, lengthscale_num_dims)),
             )
+            self.register_constraint("raw_lengthscale", lengthscale_constraint)
+
             if lengthscale_prior is not None:
                 self.register_prior(
                     "lengthscale_prior", lengthscale_prior, lambda m: m.lengthscale, lambda m, v: m._set_lengthscale(v)
                 )
 
-            self.register_constraint("raw_lengthscale", lengthscale_constraint)
-
         self.distance_module = None
         # TODO: Remove this on next official PyTorch release.
         self.__pdist_supports_batch = True

gpytorch/kernels/linear_kernel.py

@@ -57,11 +57,10 @@ def __init__(self, num_dimensions=None, offset_prior=None, variance_prior=None,
             # Remove after 1.0
             warnings.warn("The `offset_prior` argument is deprecated and no longer used.", DeprecationWarning)
         self.register_parameter(name="raw_variance", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1)))
+        self.register_constraint("raw_variance", variance_constraint)
         if variance_prior is not None:
             self.register_prior("variance_prior", variance_prior, lambda m: m.variance, lambda m, v: m._set_variance(v))
 
-        self.register_constraint("raw_variance", variance_constraint)
-
     @property
     def variance(self):
         return self.raw_variance_constraint.transform(self.raw_variance)

gpytorch/kernels/periodic_kernel.py

@@ -82,6 +82,7 @@ def __init__(self, period_length_prior=None, period_length_constraint=None, **kw
         self.register_parameter(
             name="raw_period_length", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, 1))
         )
+        self.register_constraint("raw_period_length", period_length_constraint)
 
         if period_length_prior is not None:
             self.register_prior(
@@ -91,8 +92,6 @@ def __init__(self, period_length_prior=None, period_length_constraint=None, **kw
                 lambda m, v: m._set_period_length(v),
             )
 
-        self.register_constraint("raw_period_length", period_length_constraint)
-
     @property
     def period_length(self):
         return self.raw_period_length_constraint.transform(self.raw_period_length)

gpytorch/kernels/polynomial_kernel.py

@@ -51,12 +51,11 @@ def __init__(
                 power = power.item()
 
         self.power = power
+        self.register_constraint("raw_offset", offset_constraint)
 
         if offset_prior is not None:
             self.register_prior("offset_prior", offset_prior, lambda m: m.offset, lambda m, v: m._set_offset(v))
 
-        self.register_constraint("raw_offset", offset_constraint)
-
     @property
     def offset(self) -> torch.Tensor:
         return self.raw_offset_constraint.transform(self.raw_offset)

gpytorch/kernels/scale_kernel.py

@@ -68,13 +68,13 @@ def __init__(self, base_kernel, outputscale_prior=None, outputscale_constraint=N
         self.base_kernel = base_kernel
         outputscale = torch.zeros(*self.batch_shape) if len(self.batch_shape) else torch.tensor(0.0)
         self.register_parameter(name="raw_outputscale", parameter=torch.nn.Parameter(outputscale))
+        self.register_constraint("raw_outputscale", outputscale_constraint)
+
         if outputscale_prior is not None:
             self.register_prior(
                 "outputscale_prior", outputscale_prior, lambda m: m.outputscale, lambda m, v: m._set_outputscale(v)
             )
 
-        self.register_constraint("raw_outputscale", outputscale_constraint)
-
     @property
     def outputscale(self):
         return self.raw_outputscale_constraint.transform(self.raw_outputscale)

gpytorch/likelihoods/beta_likelihood.py

@@ -43,11 +43,10 @@ def __init__(self, batch_shape=torch.Size([]), scale_prior=None, scale_constrain
             scale_constraint = Positive()
 
         self.raw_scale = torch.nn.Parameter(torch.ones(*batch_shape, 1))
+        self.register_constraint("raw_scale", scale_constraint)
         if scale_prior is not None:
             self.register_prior("scale_prior", scale_prior, lambda m: m.scale, lambda m, v: m._set_scale(v))
 
-        self.register_constraint("raw_scale", scale_constraint)
-
     @property
     def scale(self):
         return self.raw_scale_constraint.transform(self.raw_scale)

gpytorch/likelihoods/laplace_likelihood.py

@@ -29,12 +29,11 @@ def __init__(self, batch_shape=torch.Size([]), noise_prior=None, noise_constrain
             noise_constraint = Positive()
 
         self.raw_noise = torch.nn.Parameter(torch.zeros(*batch_shape, 1))
+        self.register_constraint("raw_noise", noise_constraint)
 
         if noise_prior is not None:
             self.register_prior("noise_prior", noise_prior, lambda m: m.noise, lambda m, v: m._set_noise(v))
 
-        self.register_constraint("raw_noise", noise_constraint)
-
     @property
     def noise(self):
         return self.raw_noise_constraint.transform(self.raw_noise)

gpytorch/likelihoods/noise_models.py

@@ -25,11 +25,10 @@ def __init__(self, noise_prior=None, noise_constraint=None, batch_shape=torch.Si
             noise_constraint = GreaterThan(1e-4)
 
         self.register_parameter(name="raw_noise", parameter=Parameter(torch.zeros(*batch_shape, num_tasks)))
+        self.register_constraint("raw_noise", noise_constraint)
         if noise_prior is not None:
             self.register_prior("noise_prior", noise_prior, lambda m: m.noise, lambda m, v: m._set_noise(v))
 
-        self.register_constraint("raw_noise", noise_constraint)
-
     @property
     def noise(self):
         return self.raw_noise_constraint.transform(self.raw_noise)

gpytorch/module.py

@@ -262,6 +262,20 @@ def setting_closure(module, val):
                 )
             closure = param_or_closure
 
+        if prior is not None:
+            hyperparameter_shape = closure(self).shape
+            prior_shape = prior.shape()
+
+            if prior_shape != hyperparameter_shape:
+                try:
+                    prior = prior.expand(hyperparameter_shape)
+                except RuntimeError:
+                    raise RuntimeError(
+                        "Attempting to broadcast a prior that is not broadcastable! "
+                        + f"Got parameter shape {hyperparameter_shape} "
+                        + f"but prior shape {prior_shape}!"
+                    )
+
         self.add_module(name, prior)
         self._priors[name] = (prior, closure, setting_closure)
 

gpytorch/priors/smoothed_box_prior.py

@@ -71,3 +71,9 @@ def _log_prob(self, x):
         # x = "distances from box`"
         X = ((x - self._c).abs_() - self._r).clamp(min=0)
         return (self.tails.log_prob(X) - self._M).sum(-1)
+
+    def expand(self, batch_shape):
+        batch_shape = torch.Size(batch_shape)
+        return SmoothedBoxPrior(self.a.expand(batch_shape),
+                                self.b.expand(batch_shape),
+                                self.sigma.expand(batch_shape))