Add SpanningTree properties .mode, .edge_mean #2727

fritzo · 2020-12-31T19:01:35Z

This adds:

An .edge_mean property that computes edge marginals. This has a different return shape than the usual .mean properties because SpanningTree returns a packed tensor rather than a binary array. Hence I've named it slightly differently.
A .mode property to compute a maximum-weight spanning tree. This has nearly identical implementation to the greedy sample_tree_approx() method (here using argmax rather than sampling), hence I've simply forked that code for implementation.

Tested

pyro/distributions/spanning_tree.py

fritzo added 2 commits December 31, 2020 13:27

Add a SpanningTree.edge_mean() method

ff1eabc

Implement a SpanningTree.mean() method

da27dfa

fritzo added the awaiting review label Dec 31, 2020

fritzo requested a review from martinjankowiak December 31, 2020 21:37

Convert methods to properties

7213f4b

fritzo changed the title ~~Add SpanningTree methods .mode() .edge_mean()~~ Add SpanningTree properties .mode() .edge_mean() Jan 3, 2021

fritzo changed the title ~~Add SpanningTree properties .mode() .edge_mean()~~ Add SpanningTree properties .mode(), .edge_mean() Jan 3, 2021

fritzo changed the title ~~Add SpanningTree properties .mode(), .edge_mean()~~ Add SpanningTree properties .mode, .edge_mean Jan 3, 2021

martinjankowiak previously approved these changes Jan 3, 2021

View reviewed changes

pyro/distributions/spanning_tree.py Outdated Show resolved Hide resolved

Fix typo

9b74c5a

fritzo dismissed martinjankowiak’s stale review via 9b74c5a January 3, 2021 15:58

martinjankowiak approved these changes Jan 3, 2021

View reviewed changes

martinjankowiak merged commit a5f1ae6 into dev Jan 3, 2021

fritzo deleted the spanning-tree-helpers branch September 27, 2021 14:47