7252572304

Committed 18 Dec 2023 06:59PM UTC coverage: 91.931%. Remained the same

Build # 7252572304

Build Type

Pull #3302

github

Committed by

web-flow

Commit Message

Merge fc5c5c02c into 834ff633c

Pull Request Pull Request #3302: Add tutorials using normalizing flows

Run Details

26 of 29 new or added lines in 1 file covered. (89.66%)

17 existing lines in 5 files now uncovered.

22958 of 24973 relevant lines covered (91.93%)

2.29 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

99.21

/pyro/poutine/trace_struct.py

# Copyright (c) 2017-2019 Uber Technologies, Inc.
# SPDX-License-Identifier: Apache-2.0

import sys
from collections import OrderedDict
from typing import (
    Any,
    Callable,
    Dict,
    Iterable,
    Iterator,
    List,
    Literal,
    Optional,
    Set,
    Tuple,
    Union,
)

import opt_einsum
import torch

from pyro.distributions.distribution import Distribution
from pyro.distributions.score_parts import ScoreParts
from pyro.distributions.util import scale_and_mask
from pyro.ops.packed import pack
from pyro.poutine.runtime import Message
from pyro.poutine.util import is_validation_enabled
from pyro.util import warn_if_inf, warn_if_nan


def allow_all_sites(name: str, site: Message) -> bool:
    return True


class Trace:
    """
    Graph data structure denoting the relationships amongst different pyro primitives
    in the execution trace.

    An execution trace of a Pyro program is a record of every call
    to ``pyro.sample()`` and ``pyro.param()`` in a single execution of that program.
    Traces are directed graphs whose nodes represent primitive calls or input/output,
    and whose edges represent conditional dependence relationships
    between those primitive calls. They are created and populated by ``poutine.trace``.

    Each node (or site) in a trace contains the name, input and output value of the site,
    as well as additional metadata added by inference algorithms or user annotation.
    In the case of ``pyro.sample``, the trace also includes the stochastic function
    at the site, and any observed data added by users.

    Consider the following Pyro program:

        >>> def model(x):
        ...     s = pyro.param("s", torch.tensor(0.5))
        ...     z = pyro.sample("z", dist.Normal(x, s))
        ...     return z ** 2

    We can record its execution using ``pyro.poutine.trace``
    and use the resulting data structure to compute the log-joint probability
    of all of the sample sites in the execution or extract all parameters.

        >>> trace = pyro.poutine.trace(model).get_trace(0.0)
        >>> logp = trace.log_prob_sum()
        >>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]

    We can also inspect or manipulate individual nodes in the trace.
    ``trace.nodes`` contains a ``collections.OrderedDict``
    of site names and metadata corresponding to ``x``, ``s``, ``z``, and the return value:

        >>> list(name for name in trace.nodes.keys())  # doctest: +SKIP
        ["_INPUT", "s", "z", "_RETURN"]

    Values of ``trace.nodes`` are dictionaries of node metadata:

        >>> trace.nodes["z"]  # doctest: +SKIP
        {'type': 'sample', 'name': 'z', 'is_observed': False,
         'fn': Normal(), 'value': tensor(0.6480), 'args': (), 'kwargs': {},
         'infer': {}, 'scale': 1.0, 'cond_indep_stack': (),
         'done': True, 'stop': False, 'continuation': None}

    ``'infer'`` is a dictionary of user- or algorithm-specified metadata.
    ``'args'`` and ``'kwargs'`` are the arguments passed via ``pyro.sample``
    to ``fn.__call__`` or ``fn.log_prob``.
    ``'scale'`` is used to scale the log-probability of the site when computing the log-joint.
    ``'cond_indep_stack'`` contains data structures corresponding to ``pyro.plate`` contexts
    appearing in the execution.
    ``'done'``, ``'stop'``, and ``'continuation'`` are only used by Pyro's internals.

    :param string graph_type: string specifying the kind of trace graph to construct
    """

    def __init__(self, graph_type: Literal["flat", "dense"] = "flat") -> None:
        assert graph_type in ("flat", "dense"), "{} not a valid graph type".format(
            graph_type
        )
        self.graph_type = graph_type
        self.nodes: OrderedDict[str, Message] = OrderedDict()
        self._succ: OrderedDict[str, Set[str]] = OrderedDict()
        self._pred: OrderedDict[str, Set[str]] = OrderedDict()

    def __contains__(self, name: str) -> bool:
        return name in self.nodes

    def __iter__(self) -> Iterable[str]:
        return iter(self.nodes.keys())

    def __len__(self) -> int:
        return len(self.nodes)

    @property
    def edges(self) -> Iterable[Tuple[str, str]]:
        for site, adj_nodes in self._succ.items():
            for adj_node in adj_nodes:
                yield site, adj_node

    def add_node(self, site_name: str, **kwargs: Any) -> None:
        """
        :param string site_name: the name of the site to be added

        Adds a site to the trace.

        Raises an error when attempting to add a duplicate node
        instead of silently overwriting.
        """
        if site_name in self:
            site = self.nodes[site_name]
            if site["type"] != kwargs["type"]:
                # Cannot sample or observe after a param statement.
                raise RuntimeError(
                    "{} is already in the trace as a {}".format(site_name, site["type"])
                )
            elif kwargs["type"] != "param":
                # Cannot sample after a previous sample statement.
                raise RuntimeError(
                    "Multiple {} sites named '{}'".format(kwargs["type"], site_name)
                )

        # XXX should copy in case site gets mutated, or dont bother?
        self.nodes[site_name] = kwargs  # type: ignore[assignment]
        self._pred[site_name] = set()
        self._succ[site_name] = set()

    def add_edge(self, site1: str, site2: str) -> None:
        for site in (site1, site2):
            if site not in self.nodes:
                self.add_node(site)
        self._succ[site1].add(site2)
        self._pred[site2].add(site1)

    def remove_node(self, site_name: str) -> None:
        self.nodes.pop(site_name)
        for p in self._pred[site_name]:
            self._succ[p].remove(site_name)
        for s in self._succ[site_name]:
            self._pred[s].remove(site_name)
        self._pred.pop(site_name)
        self._succ.pop(site_name)

    def predecessors(self, site_name: str) -> Set[str]:
        return self._pred[site_name]

    def successors(self, site_name: str) -> Set[str]:
        return self._succ[site_name]

    def copy(self) -> "Trace":
        """
        Makes a shallow copy of self with nodes and edges preserved.
        """
        new_tr = Trace(graph_type=self.graph_type)
        new_tr.nodes.update(self.nodes)
        new_tr._succ.update(self._succ)
        new_tr._pred.update(self._pred)
        return new_tr

    def _dfs(self, site: str, visited: Set[str]) -> Iterable[str]:
        if site in visited:
            return
        for s in self._succ[site]:
            for node in self._dfs(s, visited):
                yield node
        visited.add(site)
        yield site

    def topological_sort(self, reverse: bool = False) -> List[str]:
        """
        Return a list of nodes (site names) in topologically sorted order.

        :param bool reverse: Return the list in reverse order.
        :return: list of topologically sorted nodes (site names).
        """
        visited: Set[str] = set()
        top_sorted = []
        for s in self._succ:
            for node in self._dfs(s, visited):
                top_sorted.append(node)
        return top_sorted if reverse else list(reversed(top_sorted))

    def log_prob_sum(
        self,
        site_filter: Callable[[str, Message], bool] = allow_all_sites,
    ) -> Union[torch.Tensor, float]:
        """
        Compute the site-wise log probabilities of the trace.
        Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
        Each ``log_prob_sum`` is a scalar.
        The computation of ``log_prob_sum`` is memoized.

        :returns: total log probability.
        :rtype: torch.Tensor
        """
        result = 0.0
        for name, site in self.nodes.items():
            if site["type"] == "sample" and site_filter(name, site):
                assert isinstance(site["fn"], Distribution)
                if "log_prob_sum" in site:
                    log_p = site["log_prob_sum"]
                else:
                    try:
                        log_p = site["fn"].log_prob(
                            site["value"], *site["args"], **site["kwargs"]
                        )
                    except ValueError as e:
                        _, exc_value, traceback = sys.exc_info()
                        shapes = self.format_shapes(last_site=site["name"])
                        raise ValueError(
                            "Error while computing log_prob_sum at site '{}':\n{}\n{}\n".format(
                                name, exc_value, shapes
                            )
                        ).with_traceback(traceback) from e
                    log_p = scale_and_mask(log_p, site["scale"], site["mask"]).sum()
                    site["log_prob_sum"] = log_p
                    if is_validation_enabled():
                        warn_if_nan(log_p, "log_prob_sum at site '{}'".format(name))
                        warn_if_inf(
                            log_p,
                            "log_prob_sum at site '{}'".format(name),
                            allow_neginf=True,
                        )
                result = result + log_p  # type: ignore[assignment]
        return result

    def compute_log_prob(
        self,
        site_filter: Callable[[str, Message], bool] = allow_all_sites,
    ) -> None:
        """
        Compute the site-wise log probabilities of the trace.
        Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
        Each ``log_prob_sum`` is a scalar.
        Both computations are memoized.
        """
        for name, site in self.nodes.items():
            if site["type"] == "sample" and site_filter(name, site):
                assert isinstance(site["fn"], Distribution)
                if "log_prob" not in site:
                    try:
                        log_p = site["fn"].log_prob(
                            site["value"], *site["args"], **site["kwargs"]
                        )
                    except ValueError as e:
                        _, exc_value, traceback = sys.exc_info()
                        shapes = self.format_shapes(last_site=site["name"])
                        raise ValueError(
                            "Error while computing log_prob at site '{}':\n{}\n{}".format(
                                name, exc_value, shapes
                            )
                        ).with_traceback(traceback) from e
                    site["unscaled_log_prob"] = log_p
                    log_p = scale_and_mask(log_p, site["scale"], site["mask"])
                    site["log_prob"] = log_p
                    site["log_prob_sum"] = log_p.sum()
                    if is_validation_enabled():
                        warn_if_nan(
                            site["log_prob_sum"],
                            "log_prob_sum at site '{}'".format(name),
                        )
                        warn_if_inf(
                            site["log_prob_sum"],
                            "log_prob_sum at site '{}'".format(name),
                            allow_neginf=True,
                        )

    def compute_score_parts(self) -> None:
        """
        Compute the batched local score parts at each site of the trace.
        Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
        Each ``log_prob_sum`` is a scalar.
        All computations are memoized.
        """
        for name, site in self.nodes.items():
            if site["type"] == "sample" and "score_parts" not in site:
                assert isinstance(site["fn"], Distribution)
                # Note that ScoreParts overloads the multiplication operator
                # to correctly scale each of its three parts.
                try:
                    value = site["fn"].score_parts(
                        site["value"], *site["args"], **site["kwargs"]
                    )
                except ValueError as e:
                    _, exc_value, traceback = sys.exc_info()
                    shapes = self.format_shapes(last_site=site["name"])
                    raise ValueError(
                        "Error while computing score_parts at site '{}':\n{}\n{}".format(
                            name, exc_value, shapes
                        )
                    ).with_traceback(traceback) from e
                site["unscaled_log_prob"] = value.log_prob
                value = value.scale_and_mask(site["scale"], site["mask"])
                site["score_parts"] = value
                site["log_prob"] = value.log_prob
                site["log_prob_sum"] = value.log_prob.sum()
                if is_validation_enabled():
                    warn_if_nan(
                        site["log_prob_sum"], "log_prob_sum at site '{}'".format(name)
                    )
                    warn_if_inf(
                        site["log_prob_sum"],
                        "log_prob_sum at site '{}'".format(name),
                        allow_neginf=True,
                    )

    def detach_(self) -> None:
        """
        Detach values (in-place) at each sample site of the trace.
        """
        for _, site in self.nodes.items():
            if site["type"] == "sample":
                assert site["value"] is not None
                site["value"] = site["value"].detach()

    @property
    def observation_nodes(self) -> List[str]:
        """
        :return: a list of names of observe sites
        """
        return [
            name
            for name, node in self.nodes.items()
            if node["type"] == "sample" and node["is_observed"]
        ]

    @property
    def param_nodes(self) -> List[str]:
        """
        :return: a list of names of param sites
        """
        return [name for name, node in self.nodes.items() if node["type"] == "param"]

    @property
    def stochastic_nodes(self) -> List[str]:
        """
        :return: a list of names of sample sites
        """
        return [
            name
            for name, node in self.nodes.items()
            if node["type"] == "sample" and not node["is_observed"]
        ]

    @property
    def reparameterized_nodes(self) -> List[str]:
        """
        :return: a list of names of sample sites whose stochastic functions
            are reparameterizable primitive distributions
        """
        return [
            name
            for name, node in self.nodes.items()
            if node["type"] == "sample"
            and not node["is_observed"]
            and getattr(node["fn"], "has_rsample", False)
        ]

    @property
    def nonreparam_stochastic_nodes(self) -> List[str]:
        """
        :return: a list of names of sample sites whose stochastic functions
            are not reparameterizable primitive distributions
        """
        return list(set(self.stochastic_nodes) - set(self.reparameterized_nodes))

    def iter_stochastic_nodes(self) -> Iterator[Tuple[str, Message]]:
        """
        :return: an iterator over stochastic nodes in the trace.
        """
        for name, node in self.nodes.items():
            if node["type"] == "sample" and not node["is_observed"]:
                yield name, node

    def symbolize_dims(self, plate_to_symbol: Optional[Dict[str, str]] = None) -> None:
        """
        Assign unique symbols to all tensor dimensions.
        """
        plate_to_symbol = {} if plate_to_symbol is None else plate_to_symbol
        symbol_to_dim = {}
        for site in self.nodes.values():
            if site["type"] != "sample":
                continue

            # allocate even symbols for plate dims
            dim_to_symbol: Dict[int, str] = {}
            for frame in site["cond_indep_stack"]:
                if frame.vectorized:
                    assert frame.dim is not None
                    if frame.name in plate_to_symbol:
                        symbol = plate_to_symbol[frame.name]
                    else:
                        symbol = opt_einsum.get_symbol(2 * len(plate_to_symbol))
                        plate_to_symbol[frame.name] = symbol
                    symbol_to_dim[symbol] = frame.dim
                    dim_to_symbol[frame.dim] = symbol

            # allocate odd symbols for enum dims
            assert site["infer"] is not None
            for dim, id_ in site["infer"].get("_dim_to_id", {}).items():
                symbol = opt_einsum.get_symbol(1 + 2 * id_)
                symbol_to_dim[symbol] = dim
                dim_to_symbol[dim] = symbol
            enum_dim = site["infer"].get("_enumerate_dim")
            if enum_dim is not None:
                site["infer"]["_enumerate_symbol"] = dim_to_symbol[enum_dim]
            site["infer"]["_dim_to_symbol"] = dim_to_symbol

        self.plate_to_symbol = plate_to_symbol
        self.symbol_to_dim = symbol_to_dim

    def pack_tensors(self, plate_to_symbol: Optional[Dict[str, str]] = None) -> None:
        """
        Computes packed representations of tensors in the trace.
        This should be called after :meth:`compute_log_prob` or :meth:`compute_score_parts`.
        """
        self.symbolize_dims(plate_to_symbol)
        for site in self.nodes.values():
            if site["type"] != "sample":
                continue
            assert site["infer"] is not None
            dim_to_symbol = site["infer"]["_dim_to_symbol"]
            packed = site.setdefault("packed", {})
            try:
                packed["mask"] = pack(site["mask"], dim_to_symbol)
                if "score_parts" in site:
                    log_prob, score_function, entropy_term = site["score_parts"]
                    log_prob = pack(log_prob, dim_to_symbol)
                    score_function = pack(score_function, dim_to_symbol)
                    entropy_term = pack(entropy_term, dim_to_symbol)
                    packed["score_parts"] = ScoreParts(
                        log_prob, score_function, entropy_term
                    )
                    packed["log_prob"] = log_prob
                    packed["unscaled_log_prob"] = pack(
                        site["unscaled_log_prob"], dim_to_symbol
                    )
                elif "log_prob" in site:
                    packed["log_prob"] = pack(site["log_prob"], dim_to_symbol)
                    packed["unscaled_log_prob"] = pack(
                        site["unscaled_log_prob"], dim_to_symbol
                    )
            except ValueError as e:
                _, exc_value, traceback = sys.exc_info()
                shapes = self.format_shapes(last_site=site["name"])
                raise ValueError(
                    "Error while packing tensors at site '{}':\n  {}\n{}".format(
                        site["name"], exc_value, shapes
                    )
                ).with_traceback(traceback) from e

    def format_shapes(self, title="Trace Shapes:", last_site=None):
        """
        Returns a string showing a table of the shapes of all sites in the
        trace.
        """
        if not self.nodes:
            return title
        rows = [[title]]

        rows.append(["Param Sites:"])
        for name, site in self.nodes.items():
            if site["type"] == "param":
                rows.append([name, None] + [str(size) for size in site["value"].shape])
            if name == last_site:
                break

        rows.append(["Sample Sites:"])
        for name, site in self.nodes.items():
            if site["type"] == "sample":
                # param shape
                batch_shape = getattr(site["fn"], "batch_shape", ())
                event_shape = getattr(site["fn"], "event_shape", ())
                rows.append(
                    [name + " dist", None]
                    + [str(size) for size in batch_shape]
                    + ["|", None]
                    + [str(size) for size in event_shape]
                )

                # value shape
                event_dim = len(event_shape)
                shape = getattr(site["value"], "shape", ())
                batch_shape = shape[: len(shape) - event_dim]
                event_shape = shape[len(shape) - event_dim :]
                rows.append(
                    ["value", None]
                    + [str(size) for size in batch_shape]
                    + ["|", None]
                    + [str(size) for size in event_shape]
                )

                # log_prob shape
                if "log_prob" in site:
                    batch_shape = getattr(site["log_prob"], "shape", ())
                    rows.append(
                        ["log_prob", None]
                        + [str(size) for size in batch_shape]
                        + ["|", None]
                    )
            if name == last_site:
                break

        return _format_table(rows)


def _format_table(rows):
    """
    Formats a right justified table using None as column separator.
    """
    # compute column widths
    column_widths = [0, 0, 0]
    for row in rows:
        widths = [0, 0, 0]
        j = 0
        for cell in row:
            if cell is None:
                j += 1
            else:
                widths[j] += 1
        for j in range(3):
            column_widths[j] = max(column_widths[j], widths[j])

    # justify columns
    for i, row in enumerate(rows):
        cols = [[], [], []]
        j = 0
        for cell in row:
            if cell is None:
                j += 1
            else:
                cols[j].append(cell)
        cols = [
            [""] * (width - len(col)) + col
            if direction == "r"
            else col + [""] * (width - len(col))
            for width, col, direction in zip(column_widths, cols, "rrl")
        ]
        rows[i] = sum(cols, [])

    # compute cell widths
    cell_widths = [0] * len(rows[0])
    for row in rows:
        for j, cell in enumerate(row):
            cell_widths[j] = max(cell_widths[j], len(cell))

    # justify cells
    return "\n".join(
        " ".join(cell.rjust(width) for cell, width in zip(row, cell_widths))
        for row in rows
    )

1	# Copyright (c) 2017-2019 Uber Technologies, Inc.
2	# SPDX-License-Identifier: Apache-2.0
3
4	import sys	5✔
5	from collections import OrderedDict	5✔
6	from typing import (	5✔
7	Any,
8	Callable,
9	Dict,
10	Iterable,
11	Iterator,
12	List,
13	Literal,
14	Optional,
15	Set,
16	Tuple,
17	Union,
18	)
19
20	import opt_einsum	5✔
21	import torch	5✔
22
23	from pyro.distributions.distribution import Distribution	5✔
24	from pyro.distributions.score_parts import ScoreParts	5✔
25	from pyro.distributions.util import scale_and_mask	5✔
26	from pyro.ops.packed import pack	5✔
27	from pyro.poutine.runtime import Message	5✔
28	from pyro.poutine.util import is_validation_enabled	5✔
29	from pyro.util import warn_if_inf, warn_if_nan
30
31
32	def allow_all_sites(name: str, site: Message) -> bool:	5✔
33	return True	4✔
34
35
36	class Trace:	5✔
37	"""
38	Graph data structure denoting the relationships amongst different pyro primitives
39	in the execution trace.
40
41	An execution trace of a Pyro program is a record of every call
42	to ``pyro.sample()`` and ``pyro.param()`` in a single execution of that program.
43	Traces are directed graphs whose nodes represent primitive calls or input/output,
44	and whose edges represent conditional dependence relationships
45	between those primitive calls. They are created and populated by ``poutine.trace``.
46
47	Each node (or site) in a trace contains the name, input and output value of the site,
48	as well as additional metadata added by inference algorithms or user annotation.
49	In the case of ``pyro.sample``, the trace also includes the stochastic function
50	at the site, and any observed data added by users.
51
52	Consider the following Pyro program:
53
54	>>> def model(x):
55	... s = pyro.param("s", torch.tensor(0.5))
56	... z = pyro.sample("z", dist.Normal(x, s))
57	... return z ** 2
58
59	We can record its execution using ``pyro.poutine.trace``
60	and use the resulting data structure to compute the log-joint probability
61	of all of the sample sites in the execution or extract all parameters.
62
63	>>> trace = pyro.poutine.trace(model).get_trace(0.0)
64	>>> logp = trace.log_prob_sum()
65	>>> params = [trace.nodes[name]["value"].unconstrained() for name in trace.param_nodes]
66
67	We can also inspect or manipulate individual nodes in the trace.
68	``trace.nodes`` contains a ``collections.OrderedDict``
69	of site names and metadata corresponding to ``x``, ``s``, ``z``, and the return value:
70
71	>>> list(name for name in trace.nodes.keys()) # doctest: +SKIP
72	["_INPUT", "s", "z", "_RETURN"]
73
74	Values of ``trace.nodes`` are dictionaries of node metadata:
75
76	>>> trace.nodes["z"] # doctest: +SKIP
77	{'type': 'sample', 'name': 'z', 'is_observed': False,
78	'fn': Normal(), 'value': tensor(0.6480), 'args': (), 'kwargs': {},
79	'infer': {}, 'scale': 1.0, 'cond_indep_stack': (),
80	'done': True, 'stop': False, 'continuation': None}
81
82	``'infer'`` is a dictionary of user- or algorithm-specified metadata.
83	``'args'`` and ``'kwargs'`` are the arguments passed via ``pyro.sample``
84	to ``fn.__call__`` or ``fn.log_prob``.
85	``'scale'`` is used to scale the log-probability of the site when computing the log-joint.
86	``'cond_indep_stack'`` contains data structures corresponding to ``pyro.plate`` contexts
87	appearing in the execution.
88	``'done'``, ``'stop'``, and ``'continuation'`` are only used by Pyro's internals.
89
90	:param string graph_type: string specifying the kind of trace graph to construct
91	"""
92
93	def __init__(self, graph_type: Literal["flat", "dense"] = "flat") -> None:	5✔
94	assert graph_type in ("flat", "dense"), "{} not a valid graph type".format(	5✔
95	graph_type
96	)
97	self.graph_type = graph_type	5✔
98	self.nodes: OrderedDict[str, Message] = OrderedDict()	5✔
99	self._succ: OrderedDict[str, Set[str]] = OrderedDict()	5✔
100	self._pred: OrderedDict[str, Set[str]] = OrderedDict()	5✔
101
102	def __contains__(self, name: str) -> bool:	5✔
103	return name in self.nodes	5✔
104
105	def __iter__(self) -> Iterable[str]:	5✔
106	return iter(self.nodes.keys())	2✔
107
108	def __len__(self) -> int:	5✔
109	return len(self.nodes)	1✔
110
111	@property	5✔
112	def edges(self) -> Iterable[Tuple[str, str]]:	5✔
113	for site, adj_nodes in self._succ.items():	3✔
114	for adj_node in adj_nodes:	3✔
115	yield site, adj_node	3✔
116
117	def add_node(self, site_name: str, **kwargs: Any) -> None:	5✔
118	"""
119	:param string site_name: the name of the site to be added
120
121	Adds a site to the trace.
122
123	Raises an error when attempting to add a duplicate node
124	instead of silently overwriting.
125	"""
126	if site_name in self:	5✔
127	site = self.nodes[site_name]	4✔
128	if site["type"] != kwargs["type"]:	4✔
129	# Cannot sample or observe after a param statement.
130	raise RuntimeError(	1✔
131	"{} is already in the trace as a {}".format(site_name, site["type"])
132	)
133	elif kwargs["type"] != "param":	4✔
134	# Cannot sample after a previous sample statement.
135	raise RuntimeError(	1✔
136	"Multiple {} sites named '{}'".format(kwargs["type"], site_name)
137	)
138
139	# XXX should copy in case site gets mutated, or dont bother?
140	self.nodes[site_name] = kwargs # type: ignore[assignment]	5✔
141	self._pred[site_name] = set()	5✔
142	self._succ[site_name] = set()	5✔
143
144	def add_edge(self, site1: str, site2: str) -> None:	5✔
145	for site in (site1, site2):	3✔
146	if site not in self.nodes:	3✔
147	self.add_node(site)	3✔
148	self._succ[site1].add(site2)	3✔
149	self._pred[site2].add(site1)	3✔
150
151	def remove_node(self, site_name: str) -> None:	5✔
152	self.nodes.pop(site_name)	4✔
153	for p in self._pred[site_name]:	4✔
154	self._succ[p].remove(site_name)	1✔
155	for s in self._succ[site_name]:	4✔
156	self._pred[s].remove(site_name)	×
157	self._pred.pop(site_name)	4✔
158	self._succ.pop(site_name)	4✔
159
160	def predecessors(self, site_name: str) -> Set[str]:	5✔
161	return self._pred[site_name]	2✔
162
163	def successors(self, site_name: str) -> Set[str]:	5✔
164	return self._succ[site_name]	1✔
165
166	def copy(self) -> "Trace":	5✔
167	"""
168	Makes a shallow copy of self with nodes and edges preserved.
169	"""
170	new_tr = Trace(graph_type=self.graph_type)	4✔
171	new_tr.nodes.update(self.nodes)	4✔
172	new_tr._succ.update(self._succ)	4✔
173	new_tr._pred.update(self._pred)	4✔
174	return new_tr	4✔
175
176	def _dfs(self, site: str, visited: Set[str]) -> Iterable[str]:	5✔
177	if site in visited:	3✔
178	return	3✔
179	for s in self._succ[site]:	3✔
180	for node in self._dfs(s, visited):	3✔
181	yield node	3✔
182	visited.add(site)	3✔
183	yield site	3✔
184
185	def topological_sort(self, reverse: bool = False) -> List[str]:	5✔
186	"""
187	Return a list of nodes (site names) in topologically sorted order.
188
189	:param bool reverse: Return the list in reverse order.
190	:return: list of topologically sorted nodes (site names).
191	"""
192	visited: Set[str] = set()	3✔
193	top_sorted = []	3✔
194	for s in self._succ:	3✔
195	for node in self._dfs(s, visited):	3✔
196	top_sorted.append(node)	3✔
197	return top_sorted if reverse else list(reversed(top_sorted))	3✔
198
199	def log_prob_sum(	5✔
200	self,
201	site_filter: Callable[[str, Message], bool] = allow_all_sites,
202	) -> Union[torch.Tensor, float]:
203	"""
204	Compute the site-wise log probabilities of the trace.
205	Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
206	Each ``log_prob_sum`` is a scalar.
207	The computation of ``log_prob_sum`` is memoized.
208
209	:returns: total log probability.
210	:rtype: torch.Tensor
211	"""
212	result = 0.0	3✔
213	for name, site in self.nodes.items():	3✔
214	if site["type"] == "sample" and site_filter(name, site):	3✔
215	assert isinstance(site["fn"], Distribution)	3✔
216	if "log_prob_sum" in site:	3✔
217	log_p = site["log_prob_sum"]	1✔
218	else:
219	try:	3✔
220	log_p = site["fn"].log_prob(	3✔
221	site["value"], site["args"], *site["kwargs"]
222	)
223	except ValueError as e:	1✔
224	_, exc_value, traceback = sys.exc_info()	1✔
225	shapes = self.format_shapes(last_site=site["name"])	1✔
226	raise ValueError(
227	"Error while computing log_prob_sum at site '{}':\n{}\n{}\n".format(
228	name, exc_value, shapes
229	)
230	).with_traceback(traceback) from e
231	log_p = scale_and_mask(log_p, site["scale"], site["mask"]).sum()	3✔
232	site["log_prob_sum"] = log_p	3✔
233	if is_validation_enabled():	3✔
234	warn_if_nan(log_p, "log_prob_sum at site '{}'".format(name))
235	warn_if_inf(
236	log_p,
237	"log_prob_sum at site '{}'".format(name),
238	allow_neginf=True,
239	)
240	result = result + log_p # type: ignore[assignment]	3✔
241	return result	3✔
242
243	def compute_log_prob(	5✔
244	self,
245	site_filter: Callable[[str, Message], bool] = allow_all_sites,
246	) -> None:
247	"""
248	Compute the site-wise log probabilities of the trace.
249	Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
250	Each ``log_prob_sum`` is a scalar.
251	Both computations are memoized.
252	"""
253	for name, site in self.nodes.items():	4✔
254	if site["type"] == "sample" and site_filter(name, site):	4✔
255	assert isinstance(site["fn"], Distribution)	4✔
256	if "log_prob" not in site:	4✔
257	try:	4✔
258	log_p = site["fn"].log_prob(	4✔
259	site["value"], site["args"], *site["kwargs"]
260	)
261	except ValueError as e:	1✔
262	_, exc_value, traceback = sys.exc_info()	1✔
263	shapes = self.format_shapes(last_site=site["name"])	1✔
264	raise ValueError(
265	"Error while computing log_prob at site '{}':\n{}\n{}".format(
266	name, exc_value, shapes
267	)
268	).with_traceback(traceback) from e
269	site["unscaled_log_prob"] = log_p	4✔
270	log_p = scale_and_mask(log_p, site["scale"], site["mask"])	4✔
271	site["log_prob"] = log_p	4✔
272	site["log_prob_sum"] = log_p.sum()	4✔
273	if is_validation_enabled():	4✔
274	warn_if_nan(
275	site["log_prob_sum"],
276	"log_prob_sum at site '{}'".format(name),
277	)
278	warn_if_inf(
279	site["log_prob_sum"],
280	"log_prob_sum at site '{}'".format(name),
281	allow_neginf=True,
282	)
283
284	def compute_score_parts(self) -> None:	5✔
285	"""
286	Compute the batched local score parts at each site of the trace.
287	Each ``log_prob`` has shape equal to the corresponding ``batch_shape``.
288	Each ``log_prob_sum`` is a scalar.
289	All computations are memoized.
290	"""
291	for name, site in self.nodes.items():	4✔
292	if site["type"] == "sample" and "score_parts" not in site:	4✔
293	assert isinstance(site["fn"], Distribution)	4✔
294	# Note that ScoreParts overloads the multiplication operator
295	# to correctly scale each of its three parts.
296	try:	4✔
297	value = site["fn"].score_parts(	4✔
298	site["value"], site["args"], *site["kwargs"]
299	)
300	except ValueError as e:	1✔
301	_, exc_value, traceback = sys.exc_info()	1✔
302	shapes = self.format_shapes(last_site=site["name"])	1✔
303	raise ValueError(
304	"Error while computing score_parts at site '{}':\n{}\n{}".format(
305	name, exc_value, shapes
306	)
307	).with_traceback(traceback) from e
308	site["unscaled_log_prob"] = value.log_prob	4✔
309	value = value.scale_and_mask(site["scale"], site["mask"])	4✔
310	site["score_parts"] = value	4✔
311	site["log_prob"] = value.log_prob	4✔
312	site["log_prob_sum"] = value.log_prob.sum()	4✔
313	if is_validation_enabled():	4✔
314	warn_if_nan(
315	site["log_prob_sum"], "log_prob_sum at site '{}'".format(name)
316	)
317	warn_if_inf(
318	site["log_prob_sum"],
319	"log_prob_sum at site '{}'".format(name),
320	allow_neginf=True,
321	)
322
323	def detach_(self) -> None:	5✔
324	"""
325	Detach values (in-place) at each sample site of the trace.
326	"""
327	for _, site in self.nodes.items():	2✔
328	if site["type"] == "sample":	2✔
329	assert site["value"] is not None	2✔
330	site["value"] = site["value"].detach()	2✔
331
332	@property	5✔
333	def observation_nodes(self) -> List[str]:	5✔
334	"""
335	:return: a list of names of observe sites
336	"""
337	return [	1✔
338	name
339	for name, node in self.nodes.items()
340	if node["type"] == "sample" and node["is_observed"]
341	]
342
343	@property	5✔
344	def param_nodes(self) -> List[str]:	5✔
345	"""
346	:return: a list of names of param sites
347	"""
348	return [name for name, node in self.nodes.items() if node["type"] == "param"]	1✔
349
350	@property	5✔
351	def stochastic_nodes(self) -> List[str]:	5✔
352	"""
353	:return: a list of names of sample sites
354	"""
355	return [	1✔
356	name
357	for name, node in self.nodes.items()
358	if node["type"] == "sample" and not node["is_observed"]
359	]
360
361	@property	5✔
362	def reparameterized_nodes(self) -> List[str]:	5✔
363	"""
364	:return: a list of names of sample sites whose stochastic functions
365	are reparameterizable primitive distributions
366	"""
367	return [	1✔
368	name
369	for name, node in self.nodes.items()
370	if node["type"] == "sample"
371	and not node["is_observed"]
372	and getattr(node["fn"], "has_rsample", False)
373	]
374
375	@property	5✔
376	def nonreparam_stochastic_nodes(self) -> List[str]:	5✔
377	"""
378	:return: a list of names of sample sites whose stochastic functions
379	are not reparameterizable primitive distributions
380	"""
381	return list(set(self.stochastic_nodes) - set(self.reparameterized_nodes))	1✔
382
383	def iter_stochastic_nodes(self) -> Iterator[Tuple[str, Message]]:	5✔
384	"""
385	:return: an iterator over stochastic nodes in the trace.
386	"""
387	for name, node in self.nodes.items():	4✔
388	if node["type"] == "sample" and not node["is_observed"]:	4✔
389	yield name, node	4✔
390
391	def symbolize_dims(self, plate_to_symbol: Optional[Dict[str, str]] = None) -> None:	5✔
392	"""
393	Assign unique symbols to all tensor dimensions.
394	"""
395	plate_to_symbol = {} if plate_to_symbol is None else plate_to_symbol	3✔
396	symbol_to_dim = {}	3✔
397	for site in self.nodes.values():	3✔
398	if site["type"] != "sample":	3✔
399	continue	3✔
400
401	# allocate even symbols for plate dims
402	dim_to_symbol: Dict[int, str] = {}	3✔
403	for frame in site["cond_indep_stack"]:	3✔
404	if frame.vectorized:	3✔
405	assert frame.dim is not None	3✔
406	if frame.name in plate_to_symbol:	3✔
407	symbol = plate_to_symbol[frame.name]	3✔
408	else:
409	symbol = opt_einsum.get_symbol(2 * len(plate_to_symbol))	3✔
410	plate_to_symbol[frame.name] = symbol	3✔
411	symbol_to_dim[symbol] = frame.dim	3✔
412	dim_to_symbol[frame.dim] = symbol	3✔
413
414	# allocate odd symbols for enum dims
415	assert site["infer"] is not None	3✔
416	for dim, id_ in site["infer"].get("_dim_to_id", {}).items():	3✔
417	symbol = opt_einsum.get_symbol(1 + 2 * id_)	3✔
418	symbol_to_dim[symbol] = dim	3✔
419	dim_to_symbol[dim] = symbol	3✔
420	enum_dim = site["infer"].get("_enumerate_dim")	3✔
421	if enum_dim is not None:	3✔
422	site["infer"]["_enumerate_symbol"] = dim_to_symbol[enum_dim]	3✔
423	site["infer"]["_dim_to_symbol"] = dim_to_symbol	3✔
424
425	self.plate_to_symbol = plate_to_symbol	3✔
426	self.symbol_to_dim = symbol_to_dim	3✔
427
428	def pack_tensors(self, plate_to_symbol: Optional[Dict[str, str]] = None) -> None:	5✔
429	"""
430	Computes packed representations of tensors in the trace.
431	This should be called after :meth:`compute_log_prob` or :meth:`compute_score_parts`.
432	"""
433	self.symbolize_dims(plate_to_symbol)	3✔
434	for site in self.nodes.values():	3✔
435	if site["type"] != "sample":	3✔
436	continue	3✔
437	assert site["infer"] is not None	3✔
438	dim_to_symbol = site["infer"]["_dim_to_symbol"]	3✔
439	packed = site.setdefault("packed", {})	3✔
440	try:	3✔
441	packed["mask"] = pack(site["mask"], dim_to_symbol)	3✔
442	if "score_parts" in site:	3✔
443	log_prob, score_function, entropy_term = site["score_parts"]	3✔
444	log_prob = pack(log_prob, dim_to_symbol)	3✔
445	score_function = pack(score_function, dim_to_symbol)	3✔
446	entropy_term = pack(entropy_term, dim_to_symbol)	3✔
447	packed["score_parts"] = ScoreParts(	3✔
448	log_prob, score_function, entropy_term
449	)
450	packed["log_prob"] = log_prob	3✔
451	packed["unscaled_log_prob"] = pack(	3✔
452	site["unscaled_log_prob"], dim_to_symbol
453	)
454	elif "log_prob" in site:	3✔
455	packed["log_prob"] = pack(site["log_prob"], dim_to_symbol)	3✔
456	packed["unscaled_log_prob"] = pack(	3✔
457	site["unscaled_log_prob"], dim_to_symbol
458	)
459	except ValueError as e:	1✔
460	_, exc_value, traceback = sys.exc_info()	1✔
461	shapes = self.format_shapes(last_site=site["name"])	1✔
462	raise ValueError(
463	"Error while packing tensors at site '{}':\n {}\n{}".format(
464	site["name"], exc_value, shapes
465	)
466	).with_traceback(traceback) from e
467
468	def format_shapes(self, title="Trace Shapes:", last_site=None):	5✔
469	"""
470	Returns a string showing a table of the shapes of all sites in the
471	trace.
472	"""
473	if not self.nodes:	2✔
UNCOV 474	return title	×
475	rows = [[title]]	2✔
476
477	rows.append(["Param Sites:"])	2✔
478	for name, site in self.nodes.items():	2✔
479	if site["type"] == "param":	2✔
480	rows.append([name, None] + [str(size) for size in site["value"].shape])	1✔
481	if name == last_site:	2✔
482	break	1✔
483
484	rows.append(["Sample Sites:"])	2✔
485	for name, site in self.nodes.items():	2✔
486	if site["type"] == "sample":	2✔
487	# param shape
488	batch_shape = getattr(site["fn"], "batch_shape", ())	2✔
489	event_shape = getattr(site["fn"], "event_shape", ())	2✔
490	rows.append(	2✔
491	[name + " dist", None]
492	+ [str(size) for size in batch_shape]
493	+ ["\|", None]
494	+ [str(size) for size in event_shape]
495	)
496
497	# value shape
498	event_dim = len(event_shape)	2✔
499	shape = getattr(site["value"], "shape", ())	2✔
500	batch_shape = shape[: len(shape) - event_dim]	2✔
501	event_shape = shape[len(shape) - event_dim :]	2✔
502	rows.append(	2✔
503	["value", None]
504	+ [str(size) for size in batch_shape]
505	+ ["\|", None]
506	+ [str(size) for size in event_shape]
507	)
508
509	# log_prob shape
510	if "log_prob" in site:	2✔
511	batch_shape = getattr(site["log_prob"], "shape", ())	1✔
512	rows.append(	1✔
513	["log_prob", None]
514	+ [str(size) for size in batch_shape]
515	+ ["\|", None]
516	)
517	if name == last_site:	2✔
518	break	1✔
519
520	return _format_table(rows)	2✔
521
522
523	def _format_table(rows):	5✔
524	"""
525	Formats a right justified table using None as column separator.
526	"""
527	# compute column widths
528	column_widths = [0, 0, 0]	2✔
529	for row in rows:	2✔
530	widths = [0, 0, 0]	2✔
531	j = 0	2✔
532	for cell in row:	2✔
533	if cell is None:	2✔
534	j += 1	2✔
535	else:
536	widths[j] += 1	2✔
537	for j in range(3):	2✔
538	column_widths[j] = max(column_widths[j], widths[j])	2✔
539
540	# justify columns
541	for i, row in enumerate(rows):	2✔
542	cols = [[], [], []]	2✔
543	j = 0	2✔
544	for cell in row:	2✔
545	if cell is None:	2✔
546	j += 1	2✔
547	else:
548	cols[j].append(cell)	2✔
549	cols = [	2✔
550	[""] * (width - len(col)) + col
551	if direction == "r"
552	else col + [""] * (width - len(col))
553	for width, col, direction in zip(column_widths, cols, "rrl")
554	]
555	rows[i] = sum(cols, [])	2✔
556
557	# compute cell widths
558	cell_widths = [0] * len(rows[0])	2✔
559	for row in rows:	2✔
560	for j, cell in enumerate(row):	2✔
561	cell_widths[j] = max(cell_widths[j], len(cell))	2✔
562
563	# justify cells
564	return "\n".join(	2✔
565	" ".join(cell.rjust(width) for cell, width in zip(row, cell_widths))
566	for row in rows
567	)

pyro-ppl / pyro / 7252572304

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous