# Buggy code
def expand(self, batch_shape):
    return Gamma.expand(self, batch_shape, _instance=self)


# Traceback (AttributeError)
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'

# Fixed code
def expand(self, batch_shape):
    batch_shape = torch.Size(batch_shape)
    return GammaPrior(
        self.concentration.expand(batch_shape), self.rate.expand(batch_shape)
    )