42dc18f2a13441eefdfceed905843a3d19b27072,pyro/distributions/one_one_matching.py,OneOneMatching,log_partition_function,#OneOneMatching#,100

Before Change

        d = 1 / p.sum(0)
        for _ in range(self.bp_iters):
            s = 1 / (p @ d)
            d = 1 / (s @ p)
        b = s[:, None] * d * p

        // Evaluate the Bethe free energy.
        b = b.clamp(min=finfo.tiny ** 0.5)
        b_ = (1 - b).clamp(min=finfo.tiny ** 0.5)
        free_energy = (b * (b.log() - p.log())).sum() - (b_ * b_.log()).sum()
        return shift.sum() - free_energy

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)

After Change

        // Evaluate the Bethe free energy.
        b_ = (1 - b).clamp(min=0)
        logits = logits.clamp(min=-1 / finfo.eps)
        free_energy = (b * (log(b) - logits)).sum() - (b_ * log(b_)).sum()
        log_Z = shift.sum() - free_energy
        assert torch.isfinite(log_Z)
        return log_Z

    def log_prob(self, value):
        if self._validate_args:
            self._validate_sample(value)

Instances