Repository: cornellius-gp/gpytorch

Issue:
[Bug] NUTS not working when GammaPrior in model
# 🐛 Bug

I'm trying to use the [NUTS example](https://gpytorch.readthedocs.io/en/latest/examples/01_Simple_GP_Regression/Simple_GP_Regression_Fully_Bayesian.html#Sampling-Hyperparamters-with-GPyTorch-+-NUTS) together with the model from [this example](https://gpytorch.readthedocs.io/en/latest/examples/00_Basic_Usage/Hyperparameters.html#Putting-it-Together) and get 
```
AttributeError: 'GammaPrior' object has no attribute 'concentration'
```

## To reproduce

** Code snippet to reproduce **
```python
class ExactGPWithPriors(gpytorch.models.ExactGP):
    def __init__(self, train_x, train_y, likelihood):
        super(ExactGPWithPriors, self).__init__(train_x, train_y, likelihood)
        self.mean_module = gpytorch.means.ConstantMean()

        lengthscale_prior = gpytorch.priors.GammaPrior(3.0, 6.0)
        outputscale_prior = gpytorch.priors.GammaPrior(2.0, 0.15)
            
        self.covar_module = gpytorch.kernels.ScaleKernel(
            gpytorch.kernels.RBFKernel(
                lengthscale_prior=lengthscale_prior,
            ),
            outputscale_prior=outputscale_prior
        )

        # Initialize lengthscale and outputscale to mean of priors
        self.covar_module.base_kernel.lengthscale = lengthscale_prior.mean
        self.covar_module.outputscale = outputscale_prior.mean

    def forward(self, x):
        mean_x = self.mean_module(x)
        covar_x = self.covar_module(x)
        return gpytorch.distributions.MultivariateNormal(mean_x, covar_x)

likelihood = gpytorch.likelihoods.GaussianLikelihood(
    noise_constraint=gpytorch.constraints.GreaterThan(1e-2),
)

model = ExactGPWithPriors(train_x, train_y, likelihood)

mll = gpytorch.mlls.ExactMarginalLogLikelihood(likelihood, model)

def pyro_model(x, y):
    model.pyro_sample_from_prior()
    output = model(x)
    loss = mll.pyro_factor(output, y)
    return y

nuts_kernel = pyro.infer.mcmc.NUTS(pyro_model, adapt_step_size=True)
mcmc_run = pyro.infer.mcmc.MCMC(nuts_kernel, num_samples=100, warmup_steps=200)
mcmc_run.run(train_x, train_y)
```

** Stack trace/error message **
```
Warmup:   0%|          | 0/300 [00:00, ?it/s]
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-22-1dcdc695642b> in <module>
      9 nuts_kernel = pyro.infer.mcmc.NUTS(pyro_model, adapt_step_size=True)
     10 mcmc_run = pyro.infer.mcmc.MCMC(nuts_kernel, num_samples=100, warmup_steps=200)
---> 11 mcmc_run.run(train_x, train_y)

~/.local/lib/python3.8/site-packages/pyro/poutine/messenger.py in _context_wrap(context, fn, *args, **kwargs)
      6 def _context_wrap(context, fn, *args, **kwargs):
      7     with context:
----> 8         return fn(*args, **kwargs)
      9 
     10 

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in run(self, *args, **kwargs)
    352         z_flat_acc = [[] for _ in range(self.num_chains)]
    353         with pyro.validation_enabled(not self.disable_validation):
--> 354             for x, chain_id in self.sampler.run(*args, **kwargs):
    355                 if num_samples[chain_id] == 0:
    356                     num_samples[chain_id] += 1

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in run(self, *args, **kwargs)
    161             logger = initialize_logger(logger, "", progress_bar)
    162             hook_w_logging = _add_logging_hook(logger, progress_bar, self.hook)
--> 163             for sample in _gen_samples(self.kernel, self.warmup_steps, self.num_samples, hook_w_logging,
    164                                        i if self.num_chains > 1 else None,
    165                                        *args, **kwargs):

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/api.py in _gen_samples(kernel, warmup_steps, num_samples, hook, chain_id, *args, **kwargs)
    105 
    106 def _gen_samples(kernel, warmup_steps, num_samples, hook, chain_id, *args, **kwargs):
--> 107     kernel.setup(warmup_steps, *args, **kwargs)
    108     params = kernel.initial_params
    109     # yield structure (key, value.shape) of params

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/hmc.py in setup(self, warmup_steps, *args, **kwargs)
    258         self._warmup_steps = warmup_steps
    259         if self.model is not None:
--> 260             self._initialize_model_properties(args, kwargs)
    261         potential_energy = self.potential_fn(self.initial_params)
    262         self._cache(self.initial_params, potential_energy, None)

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/hmc.py in _initialize_model_properties(self, model_args, model_kwargs)
    223 
    224     def _initialize_model_properties(self, model_args, model_kwargs):
--> 225         init_params, potential_fn, transforms, trace = initialize_model(
    226             self.model,
    227             model_args,

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/util.py in initialize_model(model, model_args, model_kwargs, transforms, max_plate_nesting, jit_compile, jit_options, skip_jit_warnings, num_chains)
    363         automatic_transform_enabled = False
    364     if max_plate_nesting is None:
--> 365         max_plate_nesting = _guess_max_plate_nesting(model, model_args, model_kwargs)
    366     # Wrap model in `poutine.enum` to enumerate over discrete latent sites.
    367     # No-op if model does not have any discrete latents.

~/.local/lib/python3.8/site-packages/pyro/infer/mcmc/util.py in _guess_max_plate_nesting(model, args, kwargs)
    231     """
    232     with poutine.block():
--> 233         model_trace = poutine.trace(model).get_trace(*args, **kwargs)
    234     sites = [site for site in model_trace.nodes.values()
    235              if site["type"] == "sample"]

~/.local/lib/python3.8/site-packages/pyro/poutine/trace_messenger.py in get_trace(self, *args, **kwargs)
    161         Calls this poutine and returns its trace instead of the function's return value.
    162         """
--> 163         self(*args, **kwargs)
    164         return self.msngr.get_trace()

~/.local/lib/python3.8/site-packages/pyro/poutine/trace_messenger.py in __call__(self, *args, **kwargs)
    141                                       args=args, kwargs=kwargs)
    142             try:
--> 143                 ret = self.fn(*args, **kwargs)
    144             except (ValueError, RuntimeError):
    145                 exc_type, exc_value, traceback = sys.exc_info()

<ipython-input-22-1dcdc695642b> in pyro_model(x, y)
      2 
      3 def pyro_model(x, y):
----> 4     model.pyro_sample_from_prior()
      5     output = model(x)
      6     loss = mll.pyro_factor(output, y)

~/.local/lib/python3.8/site-packages/gpytorch/module.py in pyro_sample_from_prior(self)
    287         parameters of the model that have GPyTorch priors registered to them.
    288         """
--> 289         return _pyro_sample_from_prior(module=self, memo=None, prefix="")
    290 
    291     def local_load_samples(self, samples_dict, memo, prefix):

~/.local/lib/python3.8/site-packages/gpytorch/module.py in _pyro_sample_from_prior(module, memo, prefix)
    396     for mname, module_ in module.named_children():
    397         submodule_prefix = prefix + ("." if prefix else "") + mname
--> 398         _pyro_sample_from_prior(module=module_, memo=memo, prefix=submodule_prefix)
    399 
    400 

~/.local/lib/python3.8/site-packages/gpytorch/module.py in _pyro_sample_from_prior(module, memo, prefix)
    391                 memo.add(prior)
    392                 prior = prior.expand(closure().shape)
--> 393                 value = pyro.sample(prefix + ("." if prefix else "") + prior_name, prior)
    394                 setting_closure(value)
    395 

~/.local/lib/python3.8/site-packages/pyro/primitives.py in sample(name, fn, *args, **kwargs)
    108             msg["is_observed"] = True
    109         # apply the stack and return its return value
--> 110         apply_stack(msg)
    111         return msg["value"]
    112 

~/.local/lib/python3.8/site-packages/pyro/poutine/runtime.py in apply_stack(initial_msg)
    193             break
    194 
--> 195     default_process_message(msg)
    196 
    197     for frame in stack[-pointer:]:

~/.local/lib/python3.8/site-packages/pyro/poutine/runtime.py in default_process_message(msg)
    154         return msg
    155 
--> 156     msg["value"] = msg["fn"](*msg["args"], **msg["kwargs"])
    157 
    158     # after fn has been called, update msg to prevent it from being called again.

~/.local/lib/python3.8/site-packages/pyro/distributions/torch_distribution.py in __call__(self, sample_shape)
     38         :rtype: torch.Tensor
     39         """
---> 40         return self.rsample(sample_shape) if self.has_rsample else self.sample(sample_shape)
     41 
     42     @property

/usr/lib/python3.8/site-packages/torch/distributions/gamma.py in rsample(self, sample_shape)
     59     def rsample(self, sample_shape=torch.Size()):
     60         shape = self._extended_shape(sample_shape)
---> 61         value = _standard_gamma(self.concentration.expand(shape)) / self.rate.expand(shape)
     62         value.detach().clamp_(min=torch.finfo(value.dtype).tiny)  # do not record in autograd graph
     63         return value

/usr/lib/python3.8/site-packages/torch/nn/modules/module.py in __getattr__(self, name)
    582             if name in modules:
    583                 return modules[name]
--> 584         raise AttributeError("'{}' object has no attribute '{}'".format(
    585             type(self).__name__, name))
    586 

AttributeError: 'GammaPrior' object has no attribute 'concentration'
```

## System information
gpytorch 0.3.6 (pip install git, https://github.com/cornellius-gp/gpytorch/commit/4a1ba02d2367e4e9dd03eb1ccbfa4707da02dd08)
torch 1.3.1
pyro 1.0.0


Discussion:
This appears to be a problem with inheriting the standard torch distribution expand behavior which, upon inspection, appears to be strangely terrible:
https://github.com/pytorch/pytorch/blob/dd52f50fc85e6710f020936c1fc5f14673508350/torch/distributions/gamma.py#L50-L57

I'm putting up a PR now that defines better expand methods for all of our priors and fixes the issue.

Resolving patch:
diff --git a/gpytorch/priors/horseshoe_prior.py b/gpytorch/priors/horseshoe_prior.py
--- a/gpytorch/priors/horseshoe_prior.py
+++ b/gpytorch/priors/horseshoe_prior.py
@@ -7,7 +7,6 @@
 from torch.distributions import HalfCauchy, Normal, constraints
 from torch.nn import Module as TModule
 
-from gpytorch.distributions import Distribution
 from gpytorch.priors.prior import Prior
 
 
@@ -56,9 +55,5 @@ def rsample(self, sample_shape=torch.Size([])):
         return param_sample
 
     def expand(self, expand_shape, _instance=None):
-        new = self._get_checked_instance(HorseshoePrior)
         batch_shape = torch.Size(expand_shape)
-        new.scale = self.scale.expand(batch_shape)
-        super(Distribution, new).__init__(batch_shape)
-        new._validate_args = self._validate_args
-        return new
+        return HorseshoePrior(self.scale.expand(batch_shape))
diff --git a/gpytorch/priors/torch_priors.py b/gpytorch/priors/torch_priors.py
--- a/gpytorch/priors/torch_priors.py
+++ b/gpytorch/priors/torch_priors.py
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 
+import torch
 from torch.distributions import Gamma, LogNormal, MultivariateNormal, Normal, Uniform
 from torch.nn import Module as TModule
 
@@ -25,7 +26,8 @@ def __init__(self, loc, scale, validate_args=False, transform=None):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Normal.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return NormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))
 
 
 class LogNormalPrior(Prior, LogNormal):
@@ -39,7 +41,8 @@ def __init__(self, loc, scale, validate_args=None, transform=None):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return LogNormal.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return LogNormalPrior(self.loc.expand(batch_shape), self.scale.expand(batch_shape))
 
 
 class UniformPrior(Prior, Uniform):
@@ -53,7 +56,8 @@ def __init__(self, a, b, validate_args=None, transform=None):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Uniform.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return UniformPrior(self.low.expand(batch_shape), self.high.expand(batch_shape))
 
 
 class GammaPrior(Prior, Gamma):
@@ -71,7 +75,11 @@ def __init__(self, concentration, rate, validate_args=False, transform=None):
         self._transform = transform
 
     def expand(self, batch_shape):
-        return Gamma.expand(self, batch_shape, _instance=self)
+        batch_shape = torch.Size(batch_shape)
+        return GammaPrior(self.concentration.expand(batch_shape), self.rate.expand(batch_shape))
+
+    def __call__(self, *args, **kwargs):
+        return super(Gamma, self).__call__(*args, **kwargs)
 
 
 class MultivariateNormalPrior(Prior, MultivariateNormal):


Test patch: