Added Type hints and exceptions in kernels (#1802)

Author: patel-zeel

gpytorch/kernels/additive_structure_kernel.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from typing import Optional, Tuple
+
 from .kernel import Kernel
 
 
@@ -39,7 +41,9 @@ def is_stationary(self) -> bool:
         """
         return self.base_kernel.is_stationary
 
-    def __init__(self, base_kernel, num_dims, active_dims=None):
+    def __init__(
+        self, base_kernel: Kernel, num_dims: int, active_dims: Optional[Tuple[int, ...]] = None,
+    ):
         super(AdditiveStructureKernel, self).__init__(active_dims=active_dims)
         self.base_kernel = base_kernel
         self.num_dims = num_dims

gpytorch/kernels/arc_kernel.py

@@ -1,7 +1,7 @@
 #!/usr/bin/env python3
 
 from math import pi
-from typing import Optional
+from typing import Callable, Optional
 
 import torch
 
@@ -98,8 +98,8 @@ class ArcKernel(Kernel):
 
     def __init__(
         self,
-        base_kernel,
-        delta_func: Optional = None,
+        base_kernel: Kernel,
+        delta_func: Optional[Callable] = None,
         angle_prior: Optional[Prior] = None,
         radius_prior: Optional[Prior] = None,
         **kwargs,
@@ -122,6 +122,8 @@ def __init__(
             name="raw_angle", parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, self.last_dim)),
         )
         if angle_prior is not None:
+            if not isinstance(angle_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(angle_prior).__name__)
             self.register_prior(
                 "angle_prior", angle_prior, lambda m: m.angle, lambda m, v: m._set_angle(v),
             )
@@ -133,6 +135,8 @@ def __init__(
         )
 
         if radius_prior is not None:
+            if not isinstance(radius_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(radius_prior).__name__)
             self.register_prior(
                 "radius_prior", radius_prior, lambda m: m.radius, lambda m, v: m._set_radius(v),
             )

gpytorch/kernels/cosine_kernel.py

@@ -1,10 +1,12 @@
 #!/usr/bin/env python3
 
 import math
+from typing import Optional
 
 import torch
 
-from ..constraints import Positive
+from ..constraints import Interval, Positive
+from ..priors import Prior
 from .kernel import Kernel
 
 
@@ -56,7 +58,12 @@ class CosineKernel(Kernel):
 
     is_stationary = True
 
-    def __init__(self, period_length_prior=None, period_length_constraint=None, **kwargs):
+    def __init__(
+        self,
+        period_length_prior: Optional[Prior] = None,
+        period_length_constraint: Optional[Interval] = None,
+        **kwargs,
+    ):
         super(CosineKernel, self).__init__(**kwargs)
 
         self.register_parameter(
@@ -67,6 +74,8 @@ def __init__(self, period_length_prior=None, period_length_constraint=None, **kw
             period_length_constraint = Positive()
 
         if period_length_prior is not None:
+            if not isinstance(period_length_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(period_length_prior).__name__)
             self.register_prior(
                 "period_length_prior",
                 period_length_prior,

gpytorch/kernels/cylindrical_kernel.py

@@ -42,7 +42,7 @@ def __init__(
         self,
         num_angular_weights: int,
         radial_base_kernel: Kernel,
-        eps: Optional[int] = 1e-6,
+        eps: Optional[float] = 1e-6,
         angular_weights_prior: Optional[Prior] = None,
         angular_weights_constraint: Optional[Interval] = None,
         alpha_prior: Optional[Prior] = None,
@@ -76,15 +76,21 @@ def __init__(
         self.register_constraint("raw_beta", beta_constraint)
 
         if angular_weights_prior is not None:
+            if not isinstance(angular_weights_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(angular_weights_prior).__name__)
             self.register_prior(
                 "angular_weights_prior",
                 angular_weights_prior,
                 lambda m: m.angular_weights,
                 lambda m, v: m._set_angular_weights(v),
             )
         if alpha_prior is not None:
+            if not isinstance(alpha_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(alpha_prior).__name__)
             self.register_prior("alpha_prior", alpha_prior, lambda m: m.alpha, lambda m, v: m._set_alpha(v))
         if beta_prior is not None:
+            if not isinstance(beta_prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(beta_prior).__name__)
             self.register_prior("beta_prior", beta_prior, lambda m: m.beta, lambda m, v: m._set_beta(v))
 
     @property

gpytorch/kernels/distributional_input_kernel.py

@@ -1,5 +1,7 @@
 #!/usr/bin/env python3
 
+from typing import Callable
+
 import torch
 
 from .kernel import Kernel
@@ -22,7 +24,9 @@ class DistributionalInputKernel(Kernel):
     """
     has_lengthscale = True
 
-    def __init__(self, distance_function, **kwargs):
+    def __init__(
+        self, distance_function: Callable, **kwargs,
+    ):
         super(DistributionalInputKernel, self).__init__(**kwargs)
         if distance_function is None:
             raise NotImplementedError("DistributionalInputKernel requires a distance function.")

gpytorch/kernels/grid_interpolation_kernel.py

@@ -73,9 +73,9 @@ def __init__(
         self,
         base_kernel: Kernel,
         grid_size: Union[int, List[int]],
-        num_dims: int = None,
+        num_dims: Optional[int] = None,
         grid_bounds: Optional[Tuple[float, float]] = None,
-        active_dims: Tuple[int, ...] = None,
+        active_dims: Optional[Tuple[int, ...]] = None,
     ):
         has_initialized_grid = 0
         grid_is_dynamic = True

gpytorch/kernels/grid_kernel.py

@@ -1,6 +1,6 @@
 #!/usr/bin/env python3
 
-from typing import List
+from typing import Optional
 
 import torch
 from torch import Tensor
@@ -46,7 +46,11 @@ class GridKernel(Kernel):
     is_stationary = True
 
     def __init__(
-        self, base_kernel: Kernel, grid: List[Tensor], interpolation_mode: bool = False, active_dims: bool = None
+        self,
+        base_kernel: Kernel,
+        grid: Tensor,
+        interpolation_mode: Optional[bool] = False,
+        active_dims: Optional[bool] = None,
     ):
         if not base_kernel.is_stationary:
             raise RuntimeError("The base_kernel for GridKernel must be stationary.")

gpytorch/kernels/index_kernel.py

@@ -1,9 +1,12 @@
 #!/usr/bin/env python3
 
+from typing import Optional
+
 import torch
 
-from ..constraints import Positive
+from ..constraints import Interval, Positive
 from ..lazy import DiagLazyTensor, InterpolatedLazyTensor, PsdSumLazyTensor, RootLazyTensor
+from ..priors import Prior
 from ..utils.broadcasting import _mul_broadcast_shape
 from .kernel import Kernel
 
@@ -43,7 +46,14 @@ class IndexKernel(Kernel):
             The element-wise log of the :math:`\mathbf v` vector.
     """
 
-    def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs):
+    def __init__(
+        self,
+        num_tasks: int,
+        rank: Optional[int] = 1,
+        prior: Optional[Prior] = None,
+        var_constraint: Optional[Interval] = None,
+        **kwargs,
+    ):
         if rank > num_tasks:
             raise RuntimeError("Cannot create a task covariance matrix larger than the number of tasks")
         super().__init__(**kwargs)
@@ -56,6 +66,8 @@ def __init__(self, num_tasks, rank=1, prior=None, var_constraint=None, **kwargs)
         )
         self.register_parameter(name="raw_var", parameter=torch.nn.Parameter(torch.randn(*self.batch_shape, num_tasks)))
         if prior is not None:
+            if not isinstance(prior, Prior):
+                raise TypeError("Expected gpytorch.priors.Prior but got " + type(prior).__name__)
             self.register_prior("IndexKernelPrior", prior, lambda m: m._eval_covar_matrix())
 
         self.register_constraint("raw_var", var_constraint)

gpytorch/kernels/inducing_point_kernel.py

@@ -2,20 +2,29 @@
 
 import copy
 import math
+from typing import Optional, Tuple
 
 import torch
+from torch import Tensor
 
 from .. import settings
 from ..distributions import MultivariateNormal
 from ..lazy import DiagLazyTensor, LowRankRootAddedDiagLazyTensor, LowRankRootLazyTensor, MatmulLazyTensor, delazify
+from ..likelihoods import Likelihood
 from ..mlls import InducingPointKernelAddedLossTerm
 from ..models import exact_prediction_strategies
 from ..utils.cholesky import psd_safe_cholesky
 from .kernel import Kernel
 
 
 class InducingPointKernel(Kernel):
-    def __init__(self, base_kernel, inducing_points, likelihood, active_dims=None):
+    def __init__(
+        self,
+        base_kernel: Kernel,
+        inducing_points: Tensor,
+        likelihood: Likelihood,
+        active_dims: Optional[Tuple[int, ...]] = None,
+    ):
         super(InducingPointKernel, self).__init__(active_dims=active_dims)
         self.base_kernel = base_kernel
         self.likelihood = likelihood

gpytorch/kernels/kernel.py

@@ -3,15 +3,17 @@
 import warnings
 from abc import abstractmethod
 from copy import deepcopy
+from typing import Optional, Tuple
 
 import torch
 from torch.nn import ModuleList
 
 from .. import settings
-from ..constraints import Positive
+from ..constraints import Interval, Positive
 from ..lazy import LazyEvaluatedKernelTensor, ZeroLazyTensor, delazify, lazify
 from ..models import exact_prediction_strategies
 from ..module import Module
+from ..priors import Prior
 from ..utils.broadcasting import _mul_broadcast_shape
 
 
@@ -131,12 +133,12 @@ class Kernel(Module):
 
     def __init__(
         self,
-        ard_num_dims=None,
-        batch_shape=torch.Size([]),
-        active_dims=None,
-        lengthscale_prior=None,
-        lengthscale_constraint=None,
-        eps=1e-6,
+        ard_num_dims: Optional[int] = None,
+        batch_shape: Optional[torch.Size] = torch.Size([]),
+        active_dims: Optional[Tuple[int, ...]] = None,
+        lengthscale_prior: Optional[Prior] = None,
+        lengthscale_constraint: Optional[Interval] = None,
+        eps: Optional[float] = 1e-6,
         **kwargs,
     ):
         super(Kernel, self).__init__()
@@ -167,6 +169,8 @@ def __init__(
                 parameter=torch.nn.Parameter(torch.zeros(*self.batch_shape, 1, lengthscale_num_dims)),
             )
             if lengthscale_prior is not None:
+                if not isinstance(lengthscale_prior, Prior):
+                    raise TypeError("Expected gpytorch.priors.Prior but got " + type(lengthscale_prior).__name__)
                 self.register_prior(
                     "lengthscale_prior", lengthscale_prior, lambda m: m.lengthscale, lambda m, v: m._set_lengthscale(v)
                 )