d7148f390db12a3f7e4b43ab584648b2e4a1882d,pyro/distributions/spanning_tree.py,SpanningTree,log_partition_function,#SpanningTree#,109

Before Change

        edge_probs = (self.edge_logits - shift).exp()
        adjacency = torch.zeros(V, V, dtype=edge_probs.dtype)
        adjacency[grid[0], grid[1]] = edge_probs
        adjacency[grid[1], grid[0]] = edge_probs
        laplacian = adjacency.sum(-1).diag() - adjacency
        truncated = laplacian[:-1, :-1]
        try:

After Change

        // See https://en.wikipedia.org/wiki/Kirchhoff%27s_theorem
        V = self.num_vertices
        v1, v2 = make_complete_graph(V).unbind(0)
        logits = self.edge_logits.new_full((V, V), -math.inf)
        logits[v1, v2] = self.edge_logits
        logits[v2, v1] = self.edge_logits
        log_diag = logits.logsumexp(-1)
        // Numerically stabilize so that laplacian has 1"s on the diagonal.
        shift = 0.5 * log_diag
        laplacian = torch.eye(V) - (logits - shift - shift[:, None]).exp()
        truncated = laplacian[:-1, :-1]
        try:
            import gpytorch

Instances