# Buggy code
def _kinetic_energy(self, r):
    r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])
    if self.full_mass:
        return 0.5 * r_flat.dot(self._inverse_mass_matrix.matmul(r_flat))
    else:
        return 0.5 * self._inverse_mass_matrix.dot(r_flat**2)


# Traceback (KeyError)
Traceback (most recent call last):
File "/home/travis/build/uber/pyro/pyro/poutine/trace_struct.py", line 128, in log_prob_sum
log_p = site["log_prob_sum"]
KeyError: 'log_prob_sum'
During handling of the above exception, another exception occurred:
Traceback (most recent call last):
File "/home/travis/build/uber/pyro/examples/baseball.py", line 309, in <module>
main(args)
File "/home/travis/build/uber/pyro/examples/baseball.py", line 272, in main
.run(partially_pooled, at_bats, hits)
File "/home/travis/build/uber/pyro/pyro/infer/abstract_infer.py", line 84, in run
for tr, logit in self._traces(*args, **kwargs):
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/mcmc.py", line 51, in _traces
for trace in self._gen_samples(self.num_samples, trace):
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/mcmc.py", line 35, in _gen_samples
trace = self.kernel.sample(trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 276, in sample
direction, tree_depth, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 159, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 178, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 178, in _build_tree
direction, tree_depth-1, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 155, in _build_tree
return self._build_basetree(z, r, z_grads, log_slice, direction, energy_current)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/nuts.py", line 130, in _build_basetree
z, r, self._potential_energy, self._inverse_mass_matrix, step_size, z_grads=z_grads)
File "/home/travis/build/uber/pyro/pyro/ops/integrator.py", line 68, in single_step_velocity_verlet
z_grads, potential_energy = _potential_grad(potential_fn, z_next)
File "/home/travis/build/uber/pyro/pyro/ops/integrator.py", line 79, in _potential_grad
potential_energy = potential_fn(z)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/hmc.py", line 170, in _potential_energy
potential_energy = -self._compute_trace_log_prob(trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/hmc.py", line 154, in _compute_trace_log_prob
return self._trace_prob_evaluator.log_prob(model_trace)
File "/home/travis/build/uber/pyro/pyro/infer/mcmc/util.py", line 132, in log_prob
return model_trace.log_prob_sum()
File "/home/travis/build/uber/pyro/pyro/poutine/trace_struct.py", line 130, in log_prob_sum
log_p = site["fn"].log_prob(site["value"], *site["args"], **site["kwargs"])
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/transformed_distribution.py", line 86, in log_prob
log_prob += _sum_rightmost(self.base_dist.log_prob(y),
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/exponential.py", line 51, in log_prob
self._validate_sample(value)
File "/home/travis/virtualenv/python3.5.6/lib/python3.5/site-packages/torch/distributions/distribution.py", line 221, in _validate_sample
raise ValueError('The value argument must be within the support')

# Fixed code
def _kinetic_energy(self, r):
    # TODO: revert to `torch.dot` in pytorch==1.0
    # See: https://github.com/uber/pyro/issues/1458
    r_flat = torch.cat([r[site_name].reshape(-1) for site_name in sorted(r)])
    if self.full_mass:
        return 0.5 * (r_flat * (self._inverse_mass_matrix.matmul(r_flat))).sum()
    else:
        return 0.5 * (self._inverse_mass_matrix * (r_flat**2)).sum()