Repository: cornellius-gp/gpytorch
Issue: Fully connected layer before mean and covariance modules

I am trying to do a non-linear transformation of input before forwarding it through mean and covariance modules. Before master branch was merged to priors (commit 1f5491e3edcac6497d3370c8aaef9a9362048a3e), I can add a fully connected layer before passing the input through mean and covariance modules to learn a non-linear representation. For example, the following script worked fine. 

```
import math
import torch
import gpytorch
from matplotlib import pyplot as plt

from torch import optim
from gpytorch.kernels import RBFKernel
from gpytorch.means import ConstantMean
from gpytorch.likelihoods import GaussianLikelihood
from gpytorch.random_variables import GaussianRandomVariable

train_x = torch.linspace(0, 1, 11)
train_y = torch.sin(train_x.data * (2 * math.pi)) + torch.randn(train_x.size()) * 0.2

class ExactGPModel(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPModel, self).__init__(train_x, train_y, likelihood)
        self.mean_module = ConstantMean(constant_bounds=(-10, 10))
        self.covar_module = RBFKernel(log_lengthscale_bounds=(-5, 5))
        self.fc = torch.nn.Linear(1, 2)

    def forward(self, x):
        x_ = self.fc(x)
        mean_x = self.mean_module(x_)
        covar_x = self.covar_module(x_)
        return GaussianRandomVariable(mean_x, covar_x)

likelihood = GaussianLikelihood(log_noise_bounds=(-5, 5))
model = ExactGPModel(train_x.data, train_y.data, likelihood)

model.train()
likelihood.train()

optimizer = torch.optim.Adam([
    {'params': model.parameters()},  
], lr=0.1)

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

training_iter = 1000
for i in range(training_iter):
    optimizer.zero_grad()
    output = model(train_x)
    loss = -mll(output, train_y)
    loss.backward()
    print('Iter %d/%d - Loss: %.3f   log_lengthscale: %.3f   log_noise: %.3f' % (
        i + 1, training_iter, loss.data[0],
        model.covar_module.log_lengthscale.data[0, 0],
        model.likelihood.log_noise.data[0]
    ))
    optimizer.step()
```

However, now I get this runtime error because of the linear fc layer:
`AttributeError: 'Linear' object has no attribute '_get_prior_for'`

How can I get this to work with the latest version of gpytorch? DKL does something similar to what I am trying to do but implements AdditiveGridInducingVariationalGP and softmax likelihood, however, in my application I'd like to use ExactGP and gaussian likelihood. Is it possible to do so? 


Changed files:
test/lazy/test_interpolated_lazy_variable.py, gpytorch/module.py

diff --git a/gpytorch/module.py b/gpytorch/module.py
index 5a7c785a..53b10fc9 100644
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -28,40 +28,6 @@ class Module(nn.Module):
                 "Invalid parameter name {}. {} has no module {}".format(parameter_name, type(self).__name__, module)
             )
 
-    def _get_prior_for(self, parameter_name):
-        """
-        Get prior for parameter
-
-        parameter_name (str): parameter name
-        """
-        if "." in parameter_name:
-            module, parameter_name = self._get_module_and_name(parameter_name)
-            return module._get_prior_for(parameter_name)
-        else:
-            if parameter_name in self._parameters:
-                return self._priors.get(parameter_name)
-            else:
-                raise AttributeError(
-                    "Module {module} has no parameter {name}".format(module=type(self).__name__, name=parameter_name)
-                )
-
-    def _get_derived_prior(self, prior_name):
-        """
-        Get derived prior from prior name
-
-        prior_name (str): the name of the derived prior
-        """
-        if "." in prior_name:
-            module, prior_name = self._get_module_and_name(prior_name)
-            return module._get_derived_prior(prior_name)
-        else:
-            if prior_name in self._parameters:
-                return self._derived_priors.get(prior_name)
-            else:
-                raise AttributeError(
-                    "Module {module} has no derived prior {name}".format(module=type(self).__name__, name=prior_name)
-                )
-
     def forward(self, *inputs, **kwargs):
         raise NotImplementedError
 
@@ -92,16 +58,24 @@ class Module(nn.Module):
                     )
         return self
 
-    def named_parameter_priors(self):
+    def named_parameter_priors(self, memo=None, prefix=""):
         """
         Returns an iterator over module parameter priors, yielding the name of
         the parameter, the parameter itself, as well as the associated prior
         (excludes parameters for which no prior has been registered)
         """
-        for name, param in self.named_parameters():
-            prior = self._get_prior_for(name)
-            if prior is not None:
-                yield name, param, prior
+        if memo is None:
+            memo = set()
+        for name, parameter in self._parameters.items():
+            if name in self._priors and self._priors[name] not in memo:
+                prior = self._priors[name]
+                memo.add(prior)
+                yield prefix + ("." if prefix else "") + name, parameter, prior
+        for mname, module in self.named_children():
+            submodule_prefix = prefix + ("." if prefix else "") + mname
+            if hasattr(module, "named_parameter_priors"):
+                for name, parameter, prior in module.named_parameter_priors(memo, submodule_prefix):
+                    yield name, parameter, prior
 
     def named_derived_priors(self, memo=None, prefix=""):
         """Returns an iterator over module derived priors, yielding both the
diff --git a/test/lazy/test_interpolated_lazy_variable.py b/test/lazy/test_interpolated_lazy_variable.py
index 654fd0e1..f9201836 100644
--- a/test/lazy/test_interpolated_lazy_variable.py
+++ b/test/lazy/test_interpolated_lazy_variable.py
@@ -6,11 +6,25 @@ from __future__ import unicode_literals
 import gpytorch
 import torch
 import unittest
+import os
+import random
 from gpytorch.lazy import NonLazyVariable, InterpolatedLazyVariable
 from gpytorch.utils import approx_equal
 
 
 class TestInterpolatedLazyVariable(unittest.TestCase):
+    def setUp(self):
+        if os.getenv("UNLOCK_SEED") is None or os.getenv("UNLOCK_SEED").lower() == "false":
+            self.rng_state = torch.get_rng_state()
+            torch.manual_seed(0)
+            if torch.cuda.is_available():
+                torch.cuda.manual_seed_all(0)
+            random.seed(0)
+
+    def tearDown(self):
+        if hasattr(self, "rng_state"):
+            torch.set_rng_state(self.rng_state)
+
     def test_matmul(self):
         left_interp_indices = torch.LongTensor([[2, 3], [3, 4], [4, 5]]).repeat(3, 1)
         left_interp_values = torch.Tensor([[1, 2], [0.5, 1], [1, 3]]).repeat(3, 1)