def _compute_downstream_costs(model_trace, guide_trace, # non_reparam_nodes): # recursively compute downstream cost nodes for all sample sites in model and guide # (even though ultimately just need for non-reparameterizable sample sites) # 1. downstream costs used for rao-blackwellization # 2. model observe sites (as well as terms that arise from the model and guide having different # dependency structures) are taken care of via 'children_in_model' below topo_sort_guide_nodes = guide_trace.topological_sort(reverse=True) topo_sort_guide_nodes = [x for x in topo_sort_guide_nodes if guide_trace.nodes[x]["type"] == "sample"] ordered_guide_nodes_dict = {n: i for i, n in enumerate(topo_sort_guide_nodes)} downstream_guide_cost_nodes = {} downstream_costs = {} stacks = get_plate_stacks(model_trace) for node in topo_sort_guide_nodes: downstream_costs[node] = MultiFrameTensor((stacks[node], model_trace.nodes[node]['log_prob'] - guide_trace.nodes[node]['log_prob'])) nodes_included_in_sum = set([node]) downstream_guide_cost_nodes[node] = set([node]) # make more efficient by ordering children appropriately (higher children first) children = [(k, -ordered_guide_nodes_dict[k]) for k in guide_trace.successors(node)] sorted_children = sorted(children, key=itemgetter(1)) for child, _ in sorted_children: child_cost_nodes = downstream_guide_cost_nodes[child] downstream_guide_cost_nodes[node].update(child_cost_nodes) if nodes_included_in_sum.isdisjoint(child_cost_nodes): # avoid duplicates downstream_costs[node].add(*downstream_costs[child].items()) # XXX nodes_included_in_sum logic could be more fine-grained, possibly leading # to speed-ups in case there are many duplicates nodes_included_in_sum.update(child_cost_nodes) missing_downstream_costs = downstream_guide_cost_nodes[node] - nodes_included_in_sum # include terms we missed because we had to avoid duplicates for missing_node in missing_downstream_costs: downstream_costs[node].add((stacks[missing_node], model_trace.nodes[missing_node]['log_prob'] - guide_trace.nodes[missing_node]['log_prob'])) # finish assembling complete downstream costs # (the above computation may be missing terms from model) for site in non_reparam_nodes: children_in_model = set() for node in downstream_guide_cost_nodes[site]: children_in_model.update(model_trace.successors(node)) # remove terms accounted for above children_in_model.difference_update(downstream_guide_cost_nodes[site]) for child in children_in_model: assert (model_trace.nodes[child]["type"] == "sample") downstream_costs[site].add((stacks[child], model_trace.nodes[child]['log_prob'])) downstream_guide_cost_nodes[site].update([child]) for k in non_reparam_nodes: downstream_costs[k] = downstream_costs[k].sum_to(guide_trace.nodes[k]["cond_indep_stack"]) return downstream_costs, downstream_guide_cost_nodes
def _compute_log_r(model_trace, guide_trace): log_r = MultiFrameTensor() stacks = get_plate_stacks(model_trace) for name, model_site in model_trace.nodes.items(): if model_site["type"] == "sample": log_r_term = model_site["log_prob"] if not model_site["is_observed"]: log_r_term = log_r_term - guide_trace.nodes[name]["log_prob"] log_r.add((stacks[name], log_r_term.detach())) return log_r
def _brute_force_compute_downstream_costs( model_trace, guide_trace, # non_reparam_nodes): guide_nodes = [ x for x in guide_trace.nodes if guide_trace.nodes[x]["type"] == "sample" ] downstream_costs, downstream_guide_cost_nodes = {}, {} stacks = get_plate_stacks(model_trace) for node in guide_nodes: downstream_costs[node] = MultiFrameTensor( (stacks[node], model_trace.nodes[node]['log_prob'] - guide_trace.nodes[node]['log_prob'])) downstream_guide_cost_nodes[node] = set([node]) descendants = guide_trace.successors(node) for desc in descendants: desc_mft = MultiFrameTensor( (stacks[desc], model_trace.nodes[desc]['log_prob'] - guide_trace.nodes[desc]['log_prob'])) downstream_costs[node].add(*desc_mft.items()) downstream_guide_cost_nodes[node].update([desc]) for site in non_reparam_nodes: children_in_model = set() for node in downstream_guide_cost_nodes[site]: children_in_model.update(model_trace.successors(node)) children_in_model.difference_update(downstream_guide_cost_nodes[site]) for child in children_in_model: assert (model_trace.nodes[child]["type"] == "sample") child_mft = MultiFrameTensor( (stacks[child], model_trace.nodes[child]['log_prob'])) downstream_costs[site].add(*child_mft.items()) downstream_guide_cost_nodes[site].update([child]) for k in non_reparam_nodes: downstream_costs[k] = downstream_costs[k].sum_to( guide_trace.nodes[k]["cond_indep_stack"]) return downstream_costs, downstream_guide_cost_nodes