# Copyright Contributors to the Pyro project.
# SPDX-License-Identifier: Apache-2.0
import itertools
import os
from collections import defaultdict
from pathlib import Path
from types import SimpleNamespace
from typing import Callable, Collection, Dict, List, Optional, Union
import torch
import pyro
import pyro.distributions as dist
import pyro.poutine as poutine
from pyro.ops.provenance import ProvenanceTensor, detach_provenance, get_provenance
from pyro.poutine.messenger import Messenger
from pyro.poutine.util import site_is_subsample
try:
import graphviz
except ImportError:
graphviz = SimpleNamespace(Digraph=object) # for type hints
def is_sample_site(msg, *, include_deterministic=False):
if msg["type"] != "sample":
return False
if site_is_subsample(msg):
return False
if not include_deterministic:
# Ignore masked observations.
if msg["is_observed"] and msg["mask"] is False:
return False
# Exclude deterministic sites.
fn = msg["fn"]
while hasattr(fn, "base_dist"):
fn = fn.base_dist
if type(fn).__name__ == "Delta":
return False
return True
def site_is_deterministic(msg: dict) -> bool:
return msg["type"] == "sample" and msg["infer"].get("_deterministic", False)
class TrackProvenance(Messenger):
def __init__(self, *, include_deterministic=False):
self.include_deterministic = include_deterministic
def _pyro_post_sample(self, msg):
if self.include_deterministic and site_is_deterministic(msg):
provenance = frozenset({msg["name"]}) # track only direct dependencies
value = detach_provenance(msg["value"])
msg["value"] = ProvenanceTensor(value, provenance)
elif is_sample_site(msg, include_deterministic=self.include_deterministic):
provenance = frozenset({msg["name"]}) # track only direct dependencies
value = detach_provenance(msg["value"])
msg["value"] = ProvenanceTensor(value, provenance)
def _pyro_post_param(self, msg):
if msg["type"] == "param":
provenance = frozenset({msg["name"]}) # track only direct dependencies
value = detach_provenance(msg["value"])
msg["value"] = ProvenanceTensor(value, provenance)
[docs]@torch.enable_grad()
def get_dependencies(
model: Callable,
model_args: Optional[tuple] = None,
model_kwargs: Optional[dict] = None,
include_deterministic: bool = False,
) -> Dict[str, object]:
r"""
Infers dependency structure about a conditioned model.
This returns a nested dictionary with structure like::
{
"prior_dependencies": {
"variable1": {"variable1": set()},
"variable2": {"variable1": set(), "variable2": set()},
...
},
"posterior_dependencies": {
"variable1": {"variable1": {"plate1"}, "variable2": set()},
...
},
}
where
- `prior_dependencies` is a dict mapping downstream latent and observed
variables to dictionaries mapping upstream latent variables on which
they depend to sets of plates inducing full dependencies.
That is, included plates introduce quadratically many dependencies as
in complete-bipartite graphs, whereas excluded plates introduce only
linearly many dependencies as in independent sets of parallel edges.
Prior dependencies follow the original model order.
- `posterior_dependencies` is a similar dict, but mapping latent
variables to the latent or observed sites on which they depend in the
posterior. Posterior dependencies are reversed from the model order.
Dependencies elide ``pyro.deterministic`` sites and ``pyro.sample(...,
Delta(...))`` sites.
**Examples**
Here is a simple example with no plates. We see every node depends on
itself, and only the latent variables appear in the posterior::
def model_1():
a = pyro.sample("a", dist.Normal(0, 1))
pyro.sample("b", dist.Normal(a, 1), obs=torch.tensor(0.0))
assert get_dependencies(model_1) == {
"prior_dependencies": {
"a": {"a": set()},
"b": {"a": set(), "b": set()},
},
"posterior_dependencies": {
"a": {"a": set(), "b": set()},
},
}
Here is an example where two variables ``a`` and ``b`` start out
conditionally independent in the prior, but become conditionally dependent
in the posterior to the so-called collider variable ``c`` on which they
both depend. This is called "moralization" in the graphical model
literature::
def model_2():
a = pyro.sample("a", dist.Normal(0, 1))
b = pyro.sample("b", dist.LogNormal(0, 1))
c = pyro.sample("c", dist.Normal(a, b))
pyro.sample("d", dist.Normal(c, 1), obs=torch.tensor(0.))
assert get_dependencies(model_2) == {
"prior_dependencies": {
"a": {"a": set()},
"b": {"b": set()},
"c": {"a": set(), "b": set(), "c": set()},
"d": {"c": set(), "d": set()},
},
"posterior_dependencies": {
"a": {"a": set(), "b": set(), "c": set()},
"b": {"b": set(), "c": set()},
"c": {"c": set(), "d": set()},
},
}
Dependencies can be more complex in the presence of plates. So far all the
dict values have been empty sets of plates, but in the following posterior
we see that ``a`` depends on itself across the plate ``p``. This means
that, among the elements of ``a``, e.g. ``a[0]`` depends on ``a[1]`` (this
is why we explicitly allow variables to depend on themselves)::
def model_3():
with pyro.plate("p", 5):
a = pyro.sample("a", dist.Normal(0, 1))
pyro.sample("b", dist.Normal(a.sum(), 1), obs=torch.tensor(0.0))
assert get_dependencies(model_3) == {
"prior_dependencies": {
"a": {"a": set()},
"b": {"a": set(), "b": set()},
},
"posterior_dependencies": {
"a": {"a": {"p"}, "b": set()},
},
}
[1] S.Webb, A.Goliński, R.Zinkov, N.Siddharth, T.Rainforth, Y.W.Teh, F.Wood (2018)
"Faithful inversion of generative models for effective amortized inference"
https://dl.acm.org/doi/10.5555/3327144.3327229
:param callable model: A model.
:param tuple model_args: Optional tuple of model args.
:param dict model_kwargs: Optional dict of model kwargs.
:param bool include_deterministic: Whether to include deterministic sites.
:returns: A dictionary of metadata (see above).
:rtype: dict
"""
if model_args is None:
model_args = ()
if model_kwargs is None:
model_kwargs = {}
# Collect sites with tracked provenance.
with torch.random.fork_rng(), torch.no_grad(), pyro.validation_enabled(False):
with TrackProvenance(include_deterministic=include_deterministic):
trace = poutine.trace(model).get_trace(*model_args, **model_kwargs)
sample_sites = [msg for msg in trace.nodes.values() if is_sample_site(msg)]
# Collect observations.
observed = {msg["name"] for msg in sample_sites if msg["is_observed"]}
plates = {
msg["name"]: {f.name for f in msg["cond_indep_stack"] if f.vectorized}
for msg in sample_sites
}
# Find direct prior dependencies among latent and observed sites.
prior_dependencies = {n: {n: set()} for n in plates} # no deps yet
for i, downstream in enumerate(sample_sites):
upstreams = [
u for u in sample_sites[:i] if not u["is_observed"] if u["value"].numel()
]
if not upstreams:
continue
log_prob = downstream["fn"].log_prob(downstream["value"])
provenance = get_provenance(log_prob)
for upstream in upstreams:
u = upstream["name"]
if u in provenance:
d = downstream["name"]
prior_dependencies[d][u] = set()
# Next reverse dependencies and restrict downstream nodes to latent sites.
posterior_dependencies = {n: {} for n in plates if n not in observed}
for d, upstreams in prior_dependencies.items():
for u, p in upstreams.items():
if u not in observed:
# Note the folowing reverses:
# u is henceforth downstream and d is henceforth upstream.
posterior_dependencies[u][d] = p.copy()
# Moralize: add dependencies among latent variables in each Markov blanket.
# This assumes all latents are eventually observed, at least indirectly.
order = {msg["name"]: i for i, msg in enumerate(reversed(sample_sites))}
for d, upstreams in prior_dependencies.items():
upstreams = {u: p for u, p in upstreams.items() if u not in observed}
for u1, p1 in upstreams.items():
for u2, p2 in upstreams.items():
if order[u1] <= order[u2]:
p12 = posterior_dependencies[u2].setdefault(u1, set())
p12 |= plates[u1] & plates[u2] - plates[d]
p12 |= plates[u2] & p1
p12 |= plates[u1] & p2
return {
"prior_dependencies": prior_dependencies,
"posterior_dependencies": posterior_dependencies,
}
def get_model_relations(
model: Callable,
model_args: Optional[tuple] = None,
model_kwargs: Optional[dict] = None,
include_deterministic: bool = False,
):
"""
Infer relations of RVs and plates from given model and optionally data.
See https://github.com/pyro-ppl/pyro/issues/949 for more details.
This returns a dictionary with keys:
- "sample_sample" map each downstream sample site to a list of the upstream
sample sites on which it depend;
- "sample_dist" maps each sample site to the name of the distribution at
that site;
- "plate_sample" maps each plate name to a list of the sample sites within
that plate; and
- "observe" is a list of observed sample sites.
For example for the model::
def model(data):
m = pyro.sample('m', dist.Normal(0, 1))
sd = pyro.sample('sd', dist.LogNormal(m, 1))
with pyro.plate('N', len(data)):
pyro.sample('obs', dist.Normal(m, sd), obs=data)
the relation is::
{'sample_sample': {'m': [], 'sd': ['m'], 'obs': ['m', 'sd']},
'sample_dist': {'m': 'Normal', 'sd': 'LogNormal', 'obs': 'Normal'},
'plate_sample': {'N': ['obs']},
'observed': ['obs']}
:param callable model: A model to inspect.
:param model_args: Optional tuple of model args.
:param model_kwargs: Optional dict of model kwargs.
:param bool include_deterministic: Whether to include deterministic sites.
:rtype: dict
"""
if model_args is None:
model_args = ()
if model_kwargs is None:
model_kwargs = {}
assert isinstance(model_args, tuple)
assert isinstance(model_kwargs, dict)
with torch.random.fork_rng(), torch.no_grad(), pyro.validation_enabled(False):
with TrackProvenance(include_deterministic=include_deterministic):
trace = poutine.trace(model).get_trace(*model_args, **model_kwargs)
sample_sample = {}
sample_param = {}
sample_dist = {}
param_constraint = {}
plate_sample = defaultdict(list)
observed = []
def _get_type_from_frozenname(frozen_name):
return trace.nodes[frozen_name]["type"]
for name, site in trace.nodes.items():
if site["type"] == "param":
param_constraint[name] = str(site["kwargs"]["constraint"])
if site["type"] != "sample" or site_is_subsample(site):
continue
provenance = get_provenance(
site["fn"].log_prob(site["value"])
if not site_is_deterministic(site)
else site["fn"].base_dist.log_prob(site["value"])
)
sample_sample[name] = [
upstream
for upstream in provenance
if upstream != name and _get_type_from_frozenname(upstream) == "sample"
]
sample_param[name] = [
upstream
for upstream in provenance
if upstream != name and _get_type_from_frozenname(upstream) == "param"
]
sample_dist[name] = (
_get_dist_name(site["fn"])
if not site_is_deterministic(site)
else "Deterministic"
)
for frame in site["cond_indep_stack"]:
plate_sample[frame.name].append(name)
if site["is_observed"]:
observed.append(name)
def _resolve_plate_samples(plate_samples):
for p, pv in plate_samples.items():
pv = set(pv)
for q, qv in plate_samples.items():
qv = set(qv)
if len(pv & qv) > 0 and len(pv - qv) > 0 and len(qv - pv) > 0:
plate_samples_ = plate_samples.copy()
plate_samples_[q] = pv & qv
plate_samples_[q + "__CLONE"] = qv - pv
return _resolve_plate_samples(plate_samples_)
return plate_samples
plate_sample = _resolve_plate_samples(plate_sample)
# Normalize order of variables.
def sort_by_time(names: Collection[str]) -> List[str]:
return [name for name in trace.nodes if name in names]
sample_sample = {k: sort_by_time(v) for k, v in sample_sample.items()}
sample_param = {k: sort_by_time(v) for k, v in sample_param.items()}
plate_sample = {k: sort_by_time(v) for k, v in plate_sample.items()}
observed = sort_by_time(observed)
return {
"sample_sample": sample_sample,
"sample_param": sample_param,
"sample_dist": sample_dist,
"param_constraint": param_constraint,
"plate_sample": dict(plate_sample),
"observed": observed,
}
def _get_dist_name(fn):
while isinstance(
fn, (dist.Independent, dist.ExpandedDistribution, dist.MaskedDistribution)
):
fn = fn.base_dist
return type(fn).__name__
def generate_graph_specification(
model_relations: dict, render_params: bool = False
) -> dict:
"""
Convert model relations into data structure which can be readily
converted into a network.
"""
# group nodes by plate
plate_groups = dict(model_relations["plate_sample"])
plate_rvs = {rv for rvs in plate_groups.values() for rv in rvs}
plate_groups[None] = [
rv for rv in model_relations["sample_sample"] if rv not in plate_rvs
] # RVs which are in no plate
# get set of params
params = set()
if render_params:
for rv, params_list in model_relations["sample_param"].items():
for param in params_list:
params.add(param)
plate_groups[None].extend(params)
# retain node metadata
node_data = {}
for rv in model_relations["sample_sample"]:
node_data[rv] = {
"is_observed": rv in model_relations["observed"],
"distribution": model_relations["sample_dist"][rv],
}
if render_params:
for param, constraint in model_relations["param_constraint"].items():
node_data[param] = {
"is_observed": False,
"constraint": constraint,
"distribution": None,
}
# infer plate structure
# (when the order of plates cannot be determined from subset relations,
# it follows the order in which plates appear in trace)
plate_data = {}
for plate1, plate2 in list(itertools.combinations(plate_groups, 2)):
if plate1 is None or plate2 is None:
continue
nodes1 = set(plate_groups[plate1])
nodes2 = set(plate_groups[plate2])
if nodes1 < nodes2:
plate_data[plate1] = {"parent": plate2}
elif nodes1 >= nodes2:
plate_data[plate2] = {"parent": plate1}
elif nodes1 & nodes2:
raise NotImplementedError(
f"Overlapping non-nested plates {repr(plate1)},{repr(plate2)} "
"are not supported by render_model(). To help add support see "
"https://github.com/pyro-ppl/pyro/issues/2980"
)
for plate in plate_groups:
if plate is None:
continue
if plate not in plate_data:
plate_data[plate] = {"parent": None}
# infer RV edges
edge_list = []
for target, source_list in model_relations["sample_sample"].items():
edge_list.extend([(source, target) for source in source_list])
if render_params:
for target, source_list in model_relations["sample_param"].items():
edge_list.extend([(source, target) for source in source_list])
return {
"plate_groups": plate_groups,
"plate_data": plate_data,
"node_data": node_data,
"edge_list": edge_list,
}
def _deep_merge(things: list):
if len(things) == 1:
return things[0]
# Recurse into dicts.
if isinstance(things[0], dict):
result = {}
for thing in things:
for key, value in thing.items():
if key not in result:
result[key] = _deep_merge([t[key] for t in things])
return result
# Vote for booleans.
if isinstance(things[0], bool):
if all(x is True for x in things):
return True
if all(x is False for x in things):
return False
return None # i.e. maybe
# Otherwise choose arbitrarily.
return things[0]
def render_graph(
graph_specification: dict, render_distributions: bool = False
) -> "graphviz.Digraph":
"""
Create a graphviz object given a graph specification.
:param bool render_distributions: Show distribution of each RV in plot.
"""
try:
import graphviz # noqa: F401
except ImportError as e:
raise ImportError(
"Looks like you want to use graphviz (https://graphviz.org/) "
"to render your model. "
"You need to install `graphviz` to be able to use this feature. "
"It can be installed with `pip install graphviz`."
) from e
plate_groups = graph_specification["plate_groups"]
plate_data = graph_specification["plate_data"]
node_data = graph_specification["node_data"]
edge_list = graph_specification["edge_list"]
graph = graphviz.Digraph()
# add plates
plate_graph_dict = {
plate: graphviz.Digraph(name=f"cluster_{plate}")
for plate in plate_groups
if plate is not None
}
for plate, plate_graph in plate_graph_dict.items():
plate_graph.attr(label=plate.split("__CLONE")[0], labeljust="r", labelloc="b")
plate_graph_dict[None] = graph
# add nodes
colors = {False: "white", True: "gray", None: "gray:white"}
for plate, rv_list in plate_groups.items():
cur_graph = plate_graph_dict[plate]
for rv in rv_list:
color = colors[node_data[rv]["is_observed"]]
# For sample_nodes - ellipse
if node_data[rv]["distribution"]:
shape = "ellipse"
rv_label = rv
# For param_nodes - No shape
else:
shape = "plain"
rv_label = rv.replace("$params", "")
# use different symbol for Deterministic site
node_style = (
"filled,dashed"
if node_data[rv]["distribution"] == "Deterministic"
else "filled"
)
cur_graph.node(
rv, label=rv_label, shape=shape, style=node_style, fillcolor=color
)
# add leaf nodes first
while len(plate_data) >= 1:
for plate, data in plate_data.items():
parent_plate = data["parent"]
is_leaf = True
for plate2, data2 in plate_data.items():
if plate == data2["parent"]:
is_leaf = False
break
if is_leaf:
plate_graph_dict[parent_plate].subgraph(plate_graph_dict[plate])
plate_data.pop(plate)
break
# add edges
for source, target in edge_list:
graph.edge(source, target)
# render distributions if requested
if render_distributions:
dist_label = ""
for rv, data in node_data.items():
rv_dist = data["distribution"]
if rv_dist:
dist_label += rf"{rv} ~ {rv_dist}\l"
if "constraint" in data and data["constraint"]:
dist_label += rf"{rv} : {data['constraint']}\l"
graph.node("distribution_description_node", label=dist_label, shape="plaintext")
# return whole graph
return graph
[docs]def render_model(
model: Callable,
model_args: Optional[Union[tuple, List[tuple]]] = None,
model_kwargs: Optional[Union[dict, List[dict]]] = None,
filename: Optional[str] = None,
render_distributions: bool = False,
render_params: bool = False,
render_deterministic: bool = False,
) -> "graphviz.Digraph":
"""
Renders a model using `graphviz <https://graphviz.org>`_ .
If ``filename`` is provided, this saves an image; otherwise this draws the
graph. For example usage see the
`model rendering tutorial <https://pyro.ai/examples/model_rendering.html>`_ .
:param model: Model to render.
:param model_args: Tuple of positional arguments to pass to the model, or
list of tuples for semisupervised models.
:param model_kwargs: Dict of keyword arguments to pass to the model, or
list of dicts for semisupervised models.
:param str filename: Name of file or path to file to save rendered model in.
:param bool render_distributions: Whether to include RV distribution
annotations (and param constraints) in the plot.
:param bool render_params: Whether to show params inthe plot.
:param bool render_deterministic: Whether to include deterministic sites.
:returns: A model graph.
:rtype: graphviz.Digraph
"""
# Get model relations.
if not isinstance(model_args, list) and not isinstance(model_kwargs, list):
relations = [
get_model_relations(
model,
model_args,
model_kwargs,
include_deterministic=render_deterministic,
)
]
else: # semisupervised
if isinstance(model_args, list):
if not isinstance(model_kwargs, list):
model_kwargs = [model_kwargs] * len(model_args)
elif not isinstance(model_args, list):
model_args = [model_args] * len(model_kwargs)
assert len(model_args) == len(model_kwargs)
relations = [
get_model_relations(
model, args, kwargs, include_deterministic=render_deterministic
)
for args, kwargs in zip(model_args, model_kwargs)
]
# Get graph specifications.
graph_specs = [
generate_graph_specification(r, render_params=render_params) for r in relations
]
graph_spec = _deep_merge(graph_specs)
# Render.
graph = render_graph(graph_spec, render_distributions=render_distributions)
if filename is not None:
suffix = Path(filename).suffix[1:] # remove leading period from suffix
filepath = os.path.splitext(filename)[0]
graph.render(filepath, view=False, cleanup=True, format=suffix)
return graph
__all__ = [
"get_dependencies",
"render_model",
]