3040d26d11096c12daf417cd46912bb30f5dafab,pyro/infer/mcmc/hmc.py,HMC,_kinetic_energy,#HMC#Any#,141

Before Change

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

    def _energy(self, z, r):
        return self._kinetic_energy(r) + self.potential_fn(z)

    def _reset(self):

After Change

            if inv_mass_matrix.dim() == 1:
                energy = energy + 0.5 * inv_mass_matrix.dot(r_flat ** 2)
            else:
                energy = energy + 0.5 * inv_mass_matrix.matmul(r_flat).dot(r_flat)
        return energy

    def _energy(self, z, r):

Instances