pavi-pyro 0.1.0__py3-none-any.whl

pavi/__init__.py

@@ -0,0 +1,8 @@
+"""PAVI: plate-amortized variational inference for Pyro models."""
+
+from importlib.metadata import PackageNotFoundError, version
+
+try:
+    __version__ = version("pavi-pyro")
+except PackageNotFoundError:  # Not installed, e.g. imported from a source checkout
+    __version__ = "unknown"

pavi/core/__init__.py

@@ -0,0 +1,12 @@
+from .encoding_manager import EncodingManager, FactorizedEM, JointEM
+from .pavi_autoguide import PAVIGuide
+from .transform_scheme import CascadingScheme, NormalizingFlowScheme
+
+__all__ = [
+    "PAVIGuide",
+    "CascadingScheme",
+    "NormalizingFlowScheme",
+    "EncodingManager",
+    "JointEM",
+    "FactorizedEM",
+]

pavi/core/encoding_manager.py

@@ -0,0 +1,436 @@
+"""Encoding managers: the per-instance encodings of PAVI.
+
+In PAVI, the posterior of each instance of a site (random variable) inside a
+plate is a normalizing flow whose transforms are shared across the plate and
+conditioned on a small learnable **encoding** of the instance. The encoding
+encodes the statistics of the instance's posterior, which is much cheaper than
+storing these statistics for every instance (as in the mean-field
+approximation). Encoding managers create these encodings and give each site its
+context vector. They differ in how the encoding of an instance in several
+plates (e.g. subject ``s`` and vertex ``n``) is built:
+
+- ``JointEM``: every instance has its own encoding ``E^{s,n}``, stored in one
+  table per plate combination (size: product of the plate sizes).
+- ``FactorizedEM``: the encoding is built from one encoding per plate,
+  concatenated ``[E^s, E^n]`` or summed ``E^s + E^n`` (size: sum of the plate
+  sizes).
+
+Both managers share the same options:
+
+- ``share_across_sites``: whether sites in the same plates reuse the same
+  encodings (default) or each site gets its own tables.
+- ``sparse`` and ``device_params``: sparse gradients and storage of the tables
+  on another device (e.g. CPU memory) to save GPU memory.
+- ``expose``: restrict the manager to some sites.
+
+Encodings are ``torch.nn.Embedding`` tables stored in ``guide.encodings``, with
+stable parameter names: ``encodings.<plates>.weight`` for shared tables and
+``encodings.<site>.<plates>.weight`` for per-site tables, where ``<plates>`` is
+the names of the plates of the table joined by ``|``, from the outermost to the
+innermost plate (e.g. ``encodings.subjects|vertices.weight``).
+"""
+
+from abc import ABC, abstractmethod
+from collections.abc import Iterable, Mapping
+from typing import Literal, cast
+
+import numpy as np
+import torch
+from pyro.infer.autoguide import AutoGuide
+from pyro.nn.module import PyroModule
+
+PlateKey = str | Iterable[str]
+"""A plate, or a combination of plates: a name, names joined by ``|``, or an iterable of names."""
+
+EncodingSizeType = int | Mapping[PlateKey, int]
+"""Encoding size specification: either a single integer applied to every table,
+or a mapping from plates (or plate combinations) to encoding sizes."""
+
+
+def _to_plates(key: PlateKey) -> tuple[str, ...]:
+    """Convert a plate key (``"a"``, ``"a|b"`` or ``("a", "b")``) to a tuple of plate names."""
+    if isinstance(key, str):
+        return tuple(key.split("|"))
+    return tuple(key)
+
+
+class EncodingManager(ABC):
+    """Base class of the encoding managers.
+
+    An encoding manager stores encoding tables, i.e. ``torch.nn.Embedding``
+    modules with one row per instance of a plate (or of a plate combination),
+    and builds the context vector of each site from them. Subclasses define
+    which tables a site uses (``_site_tables``) and how they are combined
+    (``_combine``).
+
+    The manager is tightly coupled with ``PAVIGuide``: during guide
+    initialization, ``init_encodings()`` creates the tables in
+    ``guide.encodings``, and ``get_encodings()`` returns the encodings of the
+    plate instances selected by ``idx``, which condition the transforms.
+
+    Args:
+        encoding_size: The size of the encodings: an ``int`` for every table, or
+            a mapping keyed by plates (see the subclasses for the accepted keys).
+            A size of 0 disables the corresponding encodings.
+        share_across_sites: If ``True`` (default), sites in the same plates
+            reuse the same tables. If ``False``, each site gets its own tables.
+        expose: Optional list of the sites this manager handles. ``None``
+            (default) means all the sites in plates. Other sites get no
+            encodings (context of size 0).
+        sparse: If ``True``, embedding gradients are sparse: only the encodings
+            of the current mini-batch are updated, which saves GPU memory.
+            Optimize them with ``torch.optim.SparseAdam`` (``InferenceModule``
+            does it automatically).
+        device_params: Device of the tables, e.g. ``"cpu"`` to keep large tables
+            out of GPU memory while the rest of the guide runs on the GPU.
+            ``None`` (default) uses the device of the guide.
+        init_scale: Standard deviation of the Normal initialization of the
+            encodings.
+
+    Attributes:
+        encoding_size: The encoding size specification, with keys converted to
+            tuples of plate names.
+        site_tables: Maps each managed site to the list of its tables, as
+            ``(path, plates)`` pairs: the path of the table in
+            ``guide.encodings`` and the plates indexing its rows.
+    """
+
+    def __init__(
+        self,
+        encoding_size: EncodingSizeType,
+        *,
+        share_across_sites: bool = True,
+        expose: list[str] | None = None,
+        sparse: bool = False,
+        device_params: str | torch.device | None = None,
+        init_scale: float = 0.1,
+    ) -> None:
+        if isinstance(encoding_size, Mapping):
+            self.encoding_size: int | dict[tuple[str, ...], int] = {_to_plates(k): v for k, v in encoding_size.items()}
+        elif isinstance(encoding_size, int):
+            self.encoding_size = encoding_size
+        else:
+            raise TypeError(f"`encoding_size` must be an int or a mapping, got {type(encoding_size).__name__}.")
+
+        self.share_across_sites = share_across_sites
+        self.expose = expose
+        self.sparse = sparse
+        self.device_params = device_params
+        self.init_scale = init_scale
+        self.site_tables: dict[str, list[tuple[tuple[str, ...], tuple[str, ...]]]] = {}
+        self._table_sizes: dict[tuple[str, ...], int] = {}
+
+    @abstractmethod
+    def _site_tables(self, plates: tuple[str, ...]) -> list[tuple[str, ...]]:
+        """Return the plates of each table used by a site in ``plates`` (outermost plate first)."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _table_size(self, plates: tuple[str, ...], name_site: str) -> int:
+        """Return the encoding size of the table indexed by ``plates``."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _combine(self, encodings: list[torch.Tensor]) -> torch.Tensor:
+        """Combine the (broadcast) encodings of the tables of a site into its context."""
+        raise NotImplementedError
+
+    @abstractmethod
+    def _context_size(self, sizes: list[int], name_site: str) -> int:
+        """Return the context size of a site whose tables have the given sizes.
+
+        Raises:
+            ValueError: If the sizes cannot be combined.
+        """
+        raise NotImplementedError
+
+    def init_encodings(self, guide: AutoGuide, *args, **kwargs) -> None:
+        """Create the encoding tables in ``guide.encodings``.
+
+        Uses ``guide.sample_plate``, ``guide.plate_size``, ``guide.plate_dim`` and
+        ``guide.device``. Called once by ``PAVIGuide.__init__``.
+
+        Args:
+            guide: The guide onto which the tables are registered.
+            *args: Unused.
+            **kwargs: Unused.
+
+        Raises:
+            KeyError: If ``encoding_size`` has no entry for a table needed by a site.
+        """
+        device = self.device_params if self.device_params is not None else guide.device
+        guide.encodings = PyroModule(name="encodings")
+        self.site_tables = {}
+        self._table_sizes = {}
+
+        for name_site, site_plates in guide.sample_plate.items():
+            if self.expose is not None and name_site not in self.expose:
+                continue
+            # Order plates from the outermost to the innermost one, so that tables
+            # do not depend on the order of the plates in the model.
+            plates = tuple(sorted(site_plates, key=lambda p: guide.plate_dim[p]))
+
+            tables: list[tuple[tuple[str, ...], tuple[str, ...]]] = []
+            for table_plates in self._site_tables(plates):
+                size = self._table_size(table_plates, name_site)
+                if size <= 0:
+                    continue
+                key = "|".join(table_plates)
+                path: tuple[str, ...] = (key,) if self.share_across_sites else (name_site, key)
+                if path not in self._table_sizes:
+                    self._add_table(guide, path, int(np.prod([guide.plate_size[p] for p in table_plates])), size, device)
+                    self._table_sizes[path] = size
+                tables.append((path, table_plates))
+
+            if tables:
+                self._context_size([self._table_sizes[path] for path, _ in tables], name_site)  # Validate the sizes
+                self.site_tables[name_site] = tables
+
+    def _add_table(self, guide: AutoGuide, path: tuple[str, ...], num_rows: int, size: int, device) -> None:
+        """Create an embedding table and register it at ``path`` under ``guide.encodings``."""
+        parent = guide.encodings
+        for name in path[:-1]:
+            if not hasattr(parent, name):
+                setattr(parent, name, PyroModule(name=name))
+            parent = getattr(parent, name)
+
+        # PyroModule[...] builds a subclass of its argument dynamically, which type checkers cannot see.
+        embedding_cls = cast(type[torch.nn.Embedding], PyroModule[torch.nn.Embedding])
+        table = embedding_cls(num_embeddings=num_rows, embedding_dim=size, sparse=self.sparse, device=device)
+        with torch.no_grad():
+            table.weight.normal_(0.0, self.init_scale)
+        if self.device_params is not None:
+            # Read by InferenceModule._apply to keep the table on this device (a plain attribute, not a tensor or module).
+            table.device_params = self.device_params  # type: ignore[assignment]
+        setattr(parent, path[-1], table)
+
+    def _get_table(self, guide: AutoGuide, path: tuple[str, ...]) -> torch.nn.Embedding:
+        module = guide.encodings
+        for name in path:
+            module = getattr(module, name)
+        return module
+
+    def get_encodings(
+        self,
+        guide: AutoGuide,
+        name_sample: str,
+        idx: dict | None = None,
+        device: str | torch.device | None = None,
+        *args,
+        **kwargs,
+    ) -> torch.Tensor:
+        """Return the context of a site for the plate instances selected by ``idx``.
+
+        Args:
+            guide: The guide holding the encoding tables.
+            name_sample: The name of the site.
+            idx: Optional mapping from plate names to subsample indices. Plates
+                missing from ``idx`` (or mapped to ``None``) use all their instances.
+            device: Device of the output. Defaults to ``guide.device``.
+            *args: Unused.
+            **kwargs: Unused.
+
+        Returns:
+            A tensor of shape ``(batch_dims..., context_size)``, broadcastable to
+            the batch shape of the site, or an empty tensor if the site has no
+            encodings.
+        """
+        if device is None:
+            device = guide.device
+        if idx is not None and not isinstance(idx, dict):
+            raise TypeError("`idx` must be a dict mapping plate names to indices.")
+
+        tables = self.site_tables.get(name_sample)
+        if not tables:
+            return torch.tensor([], device=device)
+
+        site_plates = guide.sample_plate[name_sample]
+        len_batch_shape = max(-guide.plate_dim[p] for p in site_plates)
+        plate2idx = {p: (idx or {}).get(p) for p in site_plates}
+
+        # Batch shape of the site: plate sizes at their plate dims, 1 elsewhere.
+        batch_shape = [1] * len_batch_shape
+        for plate in site_plates:
+            indices = plate2idx[plate]
+            batch_shape[guide.plate_dim[plate]] = guide.plate_size[plate] if indices is None else len(indices)
+
+        encodings = []
+        for path, table_plates in tables:
+            table = self._get_table(guide, path)
+            encoding = self._lookup(guide, table, table_plates, plate2idx, len_batch_shape)
+            encodings.append(encoding.to(device).expand(*batch_shape, table.embedding_dim))
+        return self._combine(encodings)
+
+    @staticmethod
+    def _lookup(guide: AutoGuide, table: torch.nn.Embedding, plates: tuple[str, ...], plate2idx: dict, len_batch_shape: int) -> torch.Tensor:
+        """Look up the rows of ``table`` and reshape them to the plate dims of ``plates``."""
+        table_device = table.weight.device
+        indices = []
+        for plate in plates:
+            if plate2idx[plate] is None:
+                indices.append(torch.arange(guide.plate_size[plate], device=table_device))
+            elif isinstance(plate2idx[plate], torch.Tensor):
+                indices.append(plate2idx[plate].to(device=table_device, dtype=torch.long))
+            else:
+                indices.append(torch.as_tensor(np.asarray(plate2idx[plate]), device=table_device, dtype=torch.long))
+
+        # Row-major linearization of the multi-plate indices: plates are ordered
+        # from the outermost to the innermost, as the rows of the table.
+        flat_idx = torch.zeros((), dtype=torch.long, device=table_device)
+        stride = 1
+        for i in reversed(range(len(plates))):
+            shape = [1] * len(plates)
+            shape[i] = -1
+            flat_idx = flat_idx + indices[i].reshape(shape) * stride
+            stride *= guide.plate_size[plates[i]]
+
+        # Plates are sorted by plate dim, so the looked-up axes keep their order
+        # when the singleton dims of the other plates are inserted.
+        shape = [1] * len_batch_shape
+        for plate, plate_indices in zip(plates, indices, strict=True):
+            shape[guide.plate_dim[plate]] = len(plate_indices)
+        return table(flat_idx).reshape(*shape, table.embedding_dim)
+
+    def get_sample_encodings_size(self, guide: AutoGuide, name_sample: str) -> int:
+        """Return the context size of a site (0 if the site has no encodings).
+
+        Args:
+            guide: The guide (unused, kept for API consistency).
+            name_sample: The name of the site.
+
+        Returns:
+            The size of the last dimension of ``get_encodings(guide, name_sample)``.
+        """
+        tables = self.site_tables.get(name_sample)
+        if not tables:
+            return 0
+        return self._context_size([self._table_sizes[path] for path, _ in tables], name_sample)
+
+
+class JointEM(EncodingManager):
+    """One encoding per instance, for every plate combination.
+
+    Every instance of a site gets its own encoding, even for sites in several
+    plates: the table of a plate combination has as many rows as the product of
+    the plate sizes (e.g. subjects × vertices). This is the most flexible
+    manager, and the natural choice for sites in a single plate; for nested
+    plates, the table can become very large (see ``FactorizedEM``).
+
+    Args:
+        encoding_size: An ``int`` (the same for every plate combination) or a
+            mapping keyed by plate combinations, with plate names joined by
+            ``|`` in any order (e.g. ``{"subjects": 8, "subjects|vertices": 4}``).
+            A size of 0 disables the encodings of that combination.
+        share_across_sites: If ``True`` (default), sites in the same plates share
+            one table. If ``False``, each site gets its own table.
+        **kwargs: ``expose``, ``sparse``, ``device_params`` and ``init_scale``,
+            see ``EncodingManager``.
+
+    Example:
+        ```py
+        JointEM({"subjects": 8, "subjects|vertices": 4})
+        # sites in `subjects` use `encodings.subjects.weight` (S rows of size 8),
+        # sites in both plates use `encodings.subjects|vertices.weight` (S × N rows of size 4)
+        ```
+    """
+
+    def __init__(self, encoding_size: EncodingSizeType, *, share_across_sites: bool = True, **kwargs) -> None:
+        super().__init__(encoding_size, share_across_sites=share_across_sites, **kwargs)
+        if isinstance(self.encoding_size, dict):
+            sizes: dict[frozenset[str], int] = {}
+            for plates, size in self.encoding_size.items():
+                if frozenset(plates) in sizes:
+                    raise ValueError(f"Plate combination {'|'.join(plates)} is given twice in `encoding_size`.")
+                sizes[frozenset(plates)] = size
+            self._sizes = sizes
+
+    def _site_tables(self, plates: tuple[str, ...]) -> list[tuple[str, ...]]:
+        return [plates]
+
+    def _table_size(self, plates: tuple[str, ...], name_site: str) -> int:
+        if isinstance(self.encoding_size, int):
+            return self.encoding_size
+        try:
+            return self._sizes[frozenset(plates)]
+        except KeyError:
+            raise KeyError(
+                f"Plate combination {'|'.join(plates)!r} of site {name_site!r} is missing in `encoding_size`.",
+            ) from None
+
+    def _combine(self, encodings: list[torch.Tensor]) -> torch.Tensor:
+        return encodings[0]
+
+    def _context_size(self, sizes: list[int], name_site: str) -> int:
+        return sizes[0]
+
+
+class FactorizedEM(EncodingManager):
+    """Encodings built from one encoding per plate.
+
+    The encoding of an instance in several plates combines per-plate encodings,
+    e.g. ``[E^s, E^n]`` (``agg_type="concat"``) or ``E^s + E^n``
+    (``agg_type="sum"``) for subject ``s`` and vertex ``n``. Only the sum of the
+    plate sizes is stored, instead of their product with ``JointEM``. With
+    ``share_across_sites=True`` (default), every site of a plate reuses the same
+    plate encodings: this is the configuration of the PAVI brain application,
+    whose subject encodings are refitted for new subjects.
+
+    Args:
+        encoding_size: An ``int`` (the same for every plate) or a mapping keyed
+            by single plate names (e.g. ``{"subjects": 120, "vertices": 8}``).
+            A size of 0 disables the encodings of that plate.
+        agg_type: How per-plate encodings are combined: ``"concat"`` (default;
+            context size = sum of the plate sizes) or ``"sum"`` (all the plates
+            of a site must have the same encoding size).
+        share_across_sites: If ``True`` (default), all sites of a plate share its
+            table. If ``False``, each site gets its own table per plate.
+        **kwargs: ``expose``, ``sparse``, ``device_params`` and ``init_scale``,
+            see ``EncodingManager``.
+
+    Example:
+        ```py
+        FactorizedEM({"subjects": 120, "vertices": 8})
+        # tables `encodings.subjects.weight` (S × 120) and `encodings.vertices.weight` (N × 8);
+        # a site in both plates gets a context of size 128
+        ```
+    """
+
+    def __init__(
+        self,
+        encoding_size: EncodingSizeType,
+        *,
+        agg_type: Literal["concat", "sum"] = "concat",
+        share_across_sites: bool = True,
+        **kwargs,
+    ) -> None:
+        super().__init__(encoding_size, share_across_sites=share_across_sites, **kwargs)
+        if isinstance(self.encoding_size, dict):
+            for plates in self.encoding_size:
+                if len(plates) != 1:
+                    raise ValueError(f"FactorizedEM takes one plate per `encoding_size` key, got {'|'.join(plates)!r}.")
+        if agg_type not in ("concat", "sum"):
+            raise ValueError("`agg_type` must be 'concat' or 'sum'.")
+        self.agg_type = agg_type
+
+    def _site_tables(self, plates: tuple[str, ...]) -> list[tuple[str, ...]]:
+        return [(plate,) for plate in plates]
+
+    def _table_size(self, plates: tuple[str, ...], name_site: str) -> int:
+        if isinstance(self.encoding_size, int):
+            return self.encoding_size
+        try:
+            return self.encoding_size[plates]
+        except KeyError:
+            raise KeyError(f"Plate {plates[0]!r} of site {name_site!r} is missing in `encoding_size`.") from None
+
+    def _combine(self, encodings: list[torch.Tensor]) -> torch.Tensor:
+        if self.agg_type == "concat":
+            return torch.cat(encodings, dim=-1)
+        return torch.stack(encodings, dim=0).sum(dim=0)
+
+    def _context_size(self, sizes: list[int], name_site: str) -> int:
+        if self.agg_type == "concat":
+            return sum(sizes)
+        if len(set(sizes)) > 1:
+            raise ValueError(f"With agg_type='sum', the plates of site {name_site!r} must have the same encoding size, got {sizes}.")
+        return sizes[0]