Repository: cornellius-gp/gpytorch

Issue:
named_parameter_priors won't get priors for children of torch Modules
I think the PR #196  I did fixed things like DKL to at least run, but was a bad fix because we'll fail to get priors for parameters in gpytorch.Modules that are children of torch.nn.Modules. For example, consider AdditiveKernel, which places a bunch of Kernels in a ModuleList. Since named_parameter_priors is a recursive call that won't actually run on the ModuleList because it doesn't exist, we won't actually get priors for the component Kernels even if they are defined.

Ordinarily I'd just implement a fix, but I'm actually not sure how to deal with this one. Even if we implement an in house ModuleList extension, it seems dangerous and insidious to leave this failure case in, since the intention is that GPyTorch Modules are interoperable with torch ones. Therefore, a recursive version of named_parameter_priors seems not possible -- maybe we have to expand each child of a normal torch Module in the method directly?

@Balandat 


Discussion:
@gpleiss this same issue exists for the named variational strategies.
How about instead of calling this recursively on the object method we move to a functional approach where we in the top-level call of `named_parameter_priors` recursively call a function `_extract_named_parameter_priors` like the following:

```
def _extract_named_parameter_priors(module, memo=None, prefix=""):
    if memo is None:
        memo = set()
    if hasattr(module, "_priors"):
        for name, parameter in module._parameters.items():
            if name in module._priors and module._priors[name] not in memo:
                prior = module._priors[name]
                memo.add(prior)
                yield prefix + ("." if prefix else "") + name, parameter, prior
    for mname, module_ in module.named_children():
        submodule_prefix = prefix + ("." if prefix else "") + mname
        _extract_named_priors(module=mname, memo=memo, prefix=submodule_prefix)
```

Testing this now.

Resolving patch:
diff --git a/gpytorch/module.py b/gpytorch/module.py
--- a/gpytorch/module.py
+++ b/gpytorch/module.py
@@ -3,11 +3,13 @@
 from __future__ import print_function
 from __future__ import unicode_literals
 
-import torch
 from collections import OrderedDict
+
+import torch
 from torch import nn
-from .random_variables import RandomVariable
+
 from .lazy import LazyVariable
+from .random_variables import RandomVariable
 from .variational import VariationalStrategy
 
 
@@ -58,24 +60,13 @@ def initialize(self, **kwargs):
                     )
         return self
 
-    def named_parameter_priors(self, memo=None, prefix=""):
+    def named_parameter_priors(self):
         """
         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)
         """
-        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
+        return _extract_named_parameter_priors(module=self, memo=None, prefix="")
 
     def named_derived_priors(self, memo=None, prefix=""):
         """Returns an iterator over module derived priors, yielding both the
@@ -88,20 +79,9 @@ def named_derived_priors(self, memo=None, prefix=""):
                 to be called on the parameters.
 
         """
-        if memo is None:
-            memo = set()
-        for name, (prior, pnames, tf) in self._derived_priors.items():
-            if prior is not None and prior not in memo:
-                memo.add(prior)
-                parameters = tuple(getattr(self, pname) for pname in pnames)
-                yield prefix + ("." if prefix else "") + name, prior, parameters, tf
-        for mname, module in self.named_children():
-            submodule_prefix = prefix + ("." if prefix else "") + mname
-            if hasattr(module, "_derived_priors"):
-                for name, prior, parameters, tf in module.named_derived_priors(memo, submodule_prefix):
-                    yield name, prior, parameters, tf
+        return _extract_named_derived_priors(module=self, memo=None, prefix="")
 
-    def named_variational_strategies(self, memo=None, prefix=""):
+    def named_variational_strategies(self):
         """Returns an iterator over module variational strategies, yielding both
         the name of the variational strategy as well as the strategy itself.
 
@@ -110,17 +90,7 @@ def named_variational_strategies(self, memo=None, prefix=""):
                 strategy and the strategy
 
         """
-        if memo is None:
-            memo = set()
-        for name, strategy in self._variational_strategies.items():
-            if strategy is not None and strategy not in memo:
-                memo.add(strategy)
-                yield prefix + ("." if prefix else "") + name, strategy
-        for mname, module in self.named_children():
-            submodule_prefix = prefix + ("." if prefix else "") + mname
-            if hasattr(module, "named_variational_strategies"):
-                for name, strategy in module.named_variational_strategies(memo, submodule_prefix):
-                    yield name, strategy
+        return _extract_named_variational_strategies(module=self, memo=None, prefix="")
 
     def register_parameter(self, name, parameter, prior=None):
         """
@@ -201,3 +171,49 @@ def __call__(self, *inputs, **kwargs):
         if len(outputs) == 1:
             outputs = outputs[0]
         return outputs
+
+
+def _extract_named_parameter_priors(module, memo=None, prefix=""):
+    if memo is None:
+        memo = set()
+    if hasattr(module, "_priors"):
+        for name, parameter in module._parameters.items():
+            if name in module._priors and module._priors[name] not in memo:
+                prior = module._priors[name]
+                memo.add(prior)
+                yield prefix + ("." if prefix else "") + name, parameter, prior
+    for mname, module_ in module.named_children():
+        submodule_prefix = prefix + ("." if prefix else "") + mname
+        for name, parameter, prior in _extract_named_parameter_priors(
+            module=module_, memo=memo, prefix=submodule_prefix
+        ):
+            yield name, parameter, prior
+
+
+def _extract_named_variational_strategies(module, memo=None, prefix=""):
+    if memo is None:
+        memo = set()
+    if hasattr(module, "_variational_strategies"):
+        for name, strategy in module._variational_strategies.items():
+            if strategy is not None and strategy not in memo:
+                memo.add(strategy)
+                yield prefix + ("." if prefix else "") + name, strategy
+    for mname, module_ in module.named_children():
+        submodule_prefix = prefix + ("." if prefix else "") + mname
+        for name, strategy in _extract_named_variational_strategies(module=module_, memo=memo, prefix=submodule_prefix):
+            yield name, strategy
+
+
+def _extract_named_derived_priors(module, memo=None, prefix=""):
+    if memo is None:
+        memo = set()
+    if hasattr(module, "_derived_priors"):
+        for name, (prior, pnames, tf) in module._derived_priors.items():
+            if prior is not None and prior not in memo:
+                memo.add(prior)
+                parameters = tuple(getattr(module, pname) for pname in pnames)
+                yield prefix + ("." if prefix else "") + name, prior, parameters, tf
+    for mname, module_ in module.named_children():
+        submodule_prefix = prefix + ("." if prefix else "") + mname
+        for name, prior, parameters, tf in _extract_named_derived_priors(module_, memo=memo, prefix=submodule_prefix):
+            yield name, prior, parameters, tf


Test patch: