[Ready] Implement composed kernels (#831)

* add various kernels

* add various kernel

* make lint

* small bug at isotropy

* lint

* add tests

* fix bug NaN at sqrt

* add combination kernel

* add deriving kernels

* add tests

* make clearer error message

* use :math: :meth: :class: for docs

* comment out __add__, __mul__

* comment out more methods

* Update tests due to methods removed

* Update __init__.py

* revert behavior

* use raw docs

* make lint

pyro/contrib/gp/__init__.py

@@ -67,7 +67,7 @@ def fix_param(self, param, value=None):
     def set_mode(self, mode):
         """
         Sets ``mode`` for the module to be able to use its parameters in stochastic functions.
-        It also sets ``mode`` for submodules which belong to ``Parameterized`` class.
+        It also sets ``mode`` for submodules which belong to :class:`Parameterized` class.
 
         :param str mode: Either "model" or "guide".
         """
@@ -93,8 +93,8 @@ def get_param(self, param):
 
     def _register_param(self, param, mode="model"):
         """
-        Registers a parameter to Pyro. It can be seen as a wrapper for `pyro.param()` and
-        `pyro.sample()` calls.
+        Registers a parameter to Pyro. It can be seen as a wrapper for ``pyro.param()`` and
+        ``pyro.sample()`` calls.
 
         :param str param: Name of a parameter.
         :param str mode: Either "model" or "guide".

pyro/contrib/gp/kernels/__init__.py

@@ -4,8 +4,9 @@
 from .dot_product import DotProduct, Linear, Polynomial
 from .isotropic import (Exponential, Isotropy, Matern12, Matern32, Matern52,
                         RationalQuadratic, RBF, SquaredExponential)
-from .kernel import Kernel
-from .periodic import Cosine, Periodic, SineSquaredExponential
+from .kernel import (Combination, Exponent, Kernel, Product, Sum,
+                     Transforming, VerticalScaling, Warping)
+from .periodic import Cosine, ExpSineSquared, Periodic
 from .static import Bias, Constant, WhiteNoise
 
-# flake8: noqa
+# flake8: noqa

pyro/contrib/gp/kernels/brownian.py

@@ -11,7 +11,7 @@
 class Brownian(Kernel):
     """
     This kernel correponds to a two-sided Brownion motion (Wiener process):
-    ``k(x, z) = min(|x|,|z|)`` if ``x.z >= 0`` and ``k(x, z) = 0`` otherwise.
+    :math:`k(x, z) = \min(|x|,|z|)` if :math:`x\cdot z \ge 0` and :math:`k(x, z) = 0` otherwise.
 
     Note that the input dimension of this kernel must be 1.
 

pyro/contrib/gp/kernels/dot_product.py

@@ -9,7 +9,7 @@
 
 class DotProduct(Kernel):
     """
-    Base kernel for Linear and Polonomial kernels.
+    Base kernel for kernels which depends only on :math:`x\cdot z`.
 
     :param torch.Tensor variance: Variance parameter which plays the role of scaling.
     """
@@ -24,7 +24,7 @@ def __init__(self, input_dim, variance=None, active_dims=None, name=None):
 
     def _dot_product(self, X, Z=None, diag=False):
         """
-        Returns ``X.Z``.
+        Returns :math:`X\cdot Z`.
         """
         if Z is None:
             Z = X
@@ -45,7 +45,8 @@ class Linear(DotProduct):
     is equivalent to Linear Regression.
 
     Note that here we implement the homogeneous version. To use the inhomogeneous version,
-    consider using Polynomial kernel with `degree=1` or making a combination with a Bias kernel.
+    consider using :class:`Polynomial` kernel with ``degree=1`` or making
+    a combination with a :class:`.Bias` kernel.
     """
 
     def __init__(self, input_dim, variance=None, active_dims=None, name="Linear"):
@@ -57,8 +58,8 @@ def forward(self, X, Z=None, diag=False):
 
 
 class Polynomial(DotProduct):
-    """
-    Implementation of Polynomial kernel: ``k(x, z) = (bias + x.z)^d``.
+    r"""
+    Implementation of Polynomial kernel :math:`k(x, z) = (\text{bias} + x\cdot z)^d`.
 
     :param torch.Tensor bias: Bias parameter for this kernel. Should be positive.
     :param int degree: Degree of this polynomial.

pyro/contrib/gp/kernels/isotropic.py

@@ -18,11 +18,11 @@ def _torch_sqrt(x, eps=1e-18):
 class Isotropy(Kernel):
     """
     Base kernel for a family of isotropic covariance functions which is a
-    function of the distance ``r=|x-z|``.
+    function of the distance :math:`r:=|x-z|`.
 
     By default, the parameter ``lengthscale`` has size 1. To use the
     anisotropic version (different lengthscale for each dimension),
-    make sure that lengthscale has size equal to ``input_dim``.
+    make sure that ``lengthscale`` has size equal to ``input_dim``.
 
     :param torch.Tensor variance: Variance parameter of this kernel.
     :param torch.Tensor lengthscale: Length scale parameter of this kernel.
@@ -42,8 +42,8 @@ def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None,
         self.set_constraint("lengthscale", constraints.positive)
 
     def _square_scaled_dist(self, X, Z=None):
-        """
-        Returns ``||(X-Z)/lengthscale||^2``.
+        r"""
+        Returns :math:`\|\frac{(X-Z)}{\text{lengthscale}}\|^2`.
         """
         if Z is None:
             Z = X
@@ -62,8 +62,8 @@ def _square_scaled_dist(self, X, Z=None):
         return r2
 
     def _scaled_dist(self, X, Z=None):
-        """
-        Returns ``||(X-Z)/lengthscale||``.
+        r"""
+        Returns :math:`\|\frac{(X-Z)}{\text{lengthscale}}\|`.
         """
         return _torch_sqrt(self._square_scaled_dist(X, Z))
 
@@ -76,8 +76,8 @@ def _diag(self, X):
 
 
 class RBF(Isotropy):
-    """
-    Implementation of Radial Basis Function kernel: ``exp(-0.5 * r^2 / l^2)``.
+    r"""
+    Implementation of Radial Basis Function kernel :math:`\exp\left(-0.5 \times \frac{r^2}{l^2}\right)`.
     """
 
     def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None, name="RBF"):
@@ -94,7 +94,7 @@ def forward(self, X, Z=None, diag=False):
 
 class SquaredExponential(RBF):
     """
-    Squared Exponential is another name for RBF kernel.
+    SquaredExponential is another name for :class:`RBF` kernel.
     """
 
     def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None,
@@ -103,10 +103,11 @@ def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None,
 
 
 class RationalQuadratic(Isotropy):
-    """
-    Implementation of Rational Quadratic kernel: ``(1 + 0.5 * r^2 / alpha l^2)^(-alpha)``.
+    r"""
+    Implementation of RationalQuadratic kernel
+    :math:`\left(1 + 0.5 \times \frac{r^2}{\alpha l^2}\right)^{-\alpha}`.
 
-    :param torch.Tensor scale_mixture: Scale mixture (alpha) parameter of this kernel.
+    :param torch.Tensor scale_mixture: Scale mixture (:math:`\alpha`) parameter of this kernel.
         Should have size 1.
     """
 
@@ -130,8 +131,8 @@ def forward(self, X, Z=None, diag=False):
 
 
 class Exponential(Isotropy):
-    """
-    Implementation of Exponential kernel: `exp(-r/l)`.
+    r"""
+    Implementation of Exponential kernel :math:`\exp\left(-\frac{r}{l}\right)`.
     """
 
     def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None, name="Exponential"):
@@ -155,8 +156,9 @@ def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None,
 
 
 class Matern32(Isotropy):
-    """
-    Implementation of Matern32 kernel: ``(1 + sqrt(3) * r/l) * exp(-sqrt(3) * r/l)``.
+    r"""
+    Implementation of Matern32 kernel
+    :math:`\left(1 + \sqrt{3} \times \frac{r}{l}\right) \exp\left(-\sqrt{3} \times \frac{r}{l}\right)`.
     """
 
     def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None, name="Matern32"):
@@ -173,8 +175,10 @@ def forward(self, X, Z=None, diag=False):
 
 
 class Matern52(Isotropy):
-    """
-    Implementation of Matern52 kernel: ``(1 + sqrt(5) * r/l + 5/3 * r^2 / l^2) * exp(-sqrt(5) * r/l)``.
+    r"""
+    Implementation of Matern52 kernel
+    :math:`\left(1 + \sqrt{5} \times \frac{r}{l} + \frac{5}{3} \times \frac{r^2}{l^2}\right)
+    \exp\left(-\sqrt{5} \times \frac{r}{l}\right)`.
     """
 
     def __init__(self, input_dim, variance=None, lengthscale=None, active_dims=None, name="Matern52"):