pavi-pyro 0.1.0__tar.gz

pavi_pyro-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 Alexandre Le Bris
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

pavi_pyro-0.1.0/PKG-INFO

@@ -0,0 +1,226 @@
+Metadata-Version: 2.4
+Name: pavi-pyro
+Version: 0.1.0
+Summary: Plate-amortized variational inference (PAVI) for large hierarchical Bayesian models, built on Pyro
+Keywords: variational inference,hierarchical Bayesian models,pyro,normalizing flows,neuroimaging
+Author: Alexandre Le Bris
+Author-email: Alexandre Le Bris <alexandrelebris@gmail.com>
+License-Expression: MIT
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Operating System :: OS Independent
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.14
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Mathematics
+Classifier: Typing :: Typed
+Requires-Dist: numpy>=2.5.3
+Requires-Dist: pyro-ppl>=1.9.1
+Requires-Dist: pytorch-lightning>=2.6.5
+Requires-Dist: torch>=2.15.0
+Requires-Dist: zuko>=1.6.0
+Requires-Python: >=3.14
+Project-URL: Homepage, https://github.com/Dubris/pavi-pyro
+Project-URL: Documentation, https://dubris.github.io/pavi-pyro/
+Project-URL: Repository, https://github.com/Dubris/pavi-pyro
+Project-URL: Issues, https://github.com/Dubris/pavi-pyro/issues
+Description-Content-Type: text/markdown
+
+# pavi — plate-amortized variational inference in Pyro
+
+**Scalable Bayesian inference for hierarchical models with millions of latent variables.**
+
+[![Documentation](https://img.shields.io/badge/docs-GitHub%20Pages-blue)](https://dubris.github.io/pavi-pyro/)
+[![Python](https://img.shields.io/badge/python-%E2%89%A53.14-3776ab)](https://github.com/Dubris/pavi-pyro/blob/main/pyproject.toml)
+[![Pyro](https://img.shields.io/badge/built%20on-Pyro%20%2B%20PyTorch-ee4c2c)](https://pyro.ai)
+[![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/Dubris/pavi-pyro/blob/main/LICENSE)
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/parcellations_group_vs_individuals.jpg" width="85%" alt="Population-level and individual parcellations of the cortex into 17 functional networks, estimated with PAVI on the HCP dataset">
+  <br>
+  <em>Population-level (top) and individual (bottom) parcellations of the cortex into 17 functional networks, inferred with PAVI from HCP resting-state fMRI (<a href="https://hal.science/hal-04506561">Le Bris et al., 2024</a>).</em>
+</p>
+
+`pavi` implements **PAVI** ([Rouillard et al., 2023](#citation)) for [Pyro](https://pyro.ai), the probabilistic programming library built on PyTorch. You describe your data with a hierarchical Bayesian model; `pavi` approximates the posterior distribution of its latent variables, even when there are millions of them. In particular, a main application of the algorithm and this implementation was estimating individualized functional networks over the cortex from fMRI data (see the figure above and [Le Bris et al., 2024](https://hal.science/hal-04506561)).
+
+## Why PAVI?
+
+Hierarchical models repeat the same variable across plates[^plate], and the number of latent variables quickly explodes. In the brain application, about 1,000 subjects × 60,000 cortical vertices give **60 million** latent variables. With a classical mean-field Gaussian approximation over about 20 functional networks, this requires **more than 2 billion** variational parameters.
+
+PAVI builds the variational family with **normalizing flows**: a simple base distribution is mapped to a complex one by an invertible, learnable transform $f$. For an instance $n$ of a plate,
+
+$$z^n = f(u^n;\ \phi, E^n), \qquad u^n \sim \text{base distribution}.$$
+
+- The parameters $\phi$ of the transform are **shared** by all the instances of the plate. They capture what instances have in common and act as a group-level regularization.
+- Each instance only keeps a small **encoding** $E^n \in \mathbb{R}^m$, which encodes the statistics of its posterior (hence the name).
+
+The number of parameters still grows with the number of instances $N$, but much more slowly: for large $N$ it is reduced by a factor $O(m/d)$ compared to the mean-field approximation, where $d$ is the size of each variable and $m \ll d$. The mean-field approximation itself is the special case of a standard normal base and a diagonal affine transform. New subjects can be inferred by **refitting only their encodings**.
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/pavi_plate_amortization.png" width="45%" alt="(A) Generative model and PAVI variational family: posteriors of z_1 are parameterized by a shared phi_1 and per-instance encodings E_1^n. (B) Each posterior is a base distribution mapped by an invertible transform f.">
+</p>
+
+Learn more in [Concepts](https://dubris.github.io/pavi-pyro/concepts/).
+
+## Features
+
+- **Bring your own generative model**, implemented with Pyro and PyTorch.
+- **An off-the-shelf implementation of PAVI**, with training utilities: plate-aware mini-batch sampler, PyTorch Lightning inference module, KL annealing and logging callbacks.
+- **Many customization options**:
+  - a wide choice of parameterized invertible [transforms](https://dubris.github.io/pavi-pyro/transforms/);
+  - several strategies to manage the [encodings](https://dubris.github.io/pavi-pyro/encodings/);
+  - several strategies for the base distribution, called [schemes](https://dubris.github.io/pavi-pyro/schemes/).
+
+## Installation
+
+`pavi` requires Python ≥ 3.14. It is published on PyPI as `pavi-pyro` and imported as `pavi`:
+
+```bash
+pip install pavi-pyro  # or: uv add pavi-pyro
+```
+
+See [Getting started](https://dubris.github.io/pavi-pyro/getting-started/) to install from source.
+
+## Quickstart
+
+A three-level hierarchical model: a global mean $z_2$, the means $z_1^i$ of 5 groups (plate `P_1`), and 50 observations $X^{i,j}$ per group (plate `P_0`).
+
+**1. Write the generative model in Pyro.** The model receives the data and, optionally, the indices of a mini-batch (`idx`).
+
+```python
+import pyro
+import pyro.distributions as dist
+import torch
+
+from pavi.core import JointEM, NormalizingFlowScheme, PAVIGuide
+from pavi.distributions import affine, conditional_affine
+
+NUM_GROUPS, NUM_OBS, DIM = 5, 50, 2
+
+
+def model(data=None, idx=None):
+    idx = idx or {}
+    z_2 = pyro.sample("z_2", dist.Normal(torch.zeros(DIM), 1.0).to_event(1))
+    with pyro.plate("P_1", NUM_GROUPS, subsample=idx.get("P_1"), dim=-1):
+        z_1 = pyro.sample("z_1", dist.Normal(z_2, 5.0).to_event(1))
+        with pyro.plate("P_0", NUM_OBS, subsample=idx.get("P_0"), dim=-2):
+            pyro.sample("X", dist.Normal(z_1, 1.0).to_event(1), obs=data)
+
+
+# Simulate a dataset from the model
+pyro.set_rng_seed(0)
+X = pyro.poutine.trace(model).get_trace().nodes["X"]["value"]
+```
+
+**2. Build the PAVI guide[^guide].** Choose a transform for each latent site[^site], an encoding manager and a scheme.
+
+```python
+guide = PAVIGuide(
+    model,
+    transform_modules={
+        "z_2": [affine(scale_init=1.0)],                     # global variable
+        "z_1": [conditional_affine(hidden_layers=[32, 32])],  # plate P_1: conditioned on encodings
+    },
+    transform_scheme=NormalizingFlowScheme,
+    encoding_manager=JointEM(encoding_size={"P_1": 8}),
+    hidden_sites=["X"],  # observed sites are not inferred
+    device="cpu",
+    data=X,
+)
+```
+
+**3. Fit the guide** by maximizing the evidence lower bound (ELBO).
+
+```python
+elbo = pyro.infer.Trace_ELBO(num_particles=8)
+optimizer = torch.optim.Adam(guide.parameters(), lr=1e-2)
+for step in range(2000):
+    optimizer.zero_grad()
+    loss = elbo.differentiable_loss(model, guide, X)
+    loss.backward()
+    optimizer.step()
+```
+
+**4. Sample from the approximate posterior.**
+
+```python
+with torch.no_grad():
+    posterior_sample = guide(X)
+print(posterior_sample["z_1"])  # one sample of the 5 group means
+```
+
+The [toy-model notebook](https://dubris.github.io/pavi-pyro/examples/toymodel_hgm/) compares the result with the exact posterior.
+
+## How it works
+
+```mermaid
+flowchart LR
+    M["Pyro model<br/>with plates"] -->|"inspected once"| S["latent sites<br/>and their plates"]
+    S --> EM["encoding manager<br/>encodings E"]
+    S --> T["transforms<br/>per site"]
+    EM -->|"conditions"| T
+    T --> SC["scheme<br/>base distribution"]
+    SC --> Q["approximate posterior q"]
+```
+
+When it is created, the guide runs the model once to list its latent sites and their plates. The encoding manager then creates the encodings, and each site gets its chain of transforms, conditioned on the encodings when the site lives in a plate. The scheme chooses the base distribution that the transforms are applied to. See [The PAVI guide](https://dubris.github.io/pavi-pyro/guide/).
+
+## Application: individual brain functional networks
+
+PAVI makes it possible to map the **functional networks of each individual brain** at the resolution of the cortical mesh, from resting-state fMRI ([Le Bris et al., 2024](https://hal.science/hal-04506561)). A hierarchical model explains the functional connectivity of each vertex as a mixture of 17 networks, whose spatial organization varies across individuals around a population-level atlas.
+
+Instead of a discrete atlas, the model yields a **continuous, probabilistic map of each network**, at the population level (below) and for each individual:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/probabilistic_networks_1.jpg" width="48%" alt="Population-level probabilistic maps of functional networks VIN-1 to CON-2">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/probabilistic_networks_2.jpg" width="48%" alt="Population-level probabilistic maps of functional networks CON-3 to DMN-4">
+</p>
+
+These individual maps are reliable, reflect individual-specific organization, predict behavior, and transfer to small datasets through the population-level priors. See [the brain application example](https://dubris.github.io/pavi-pyro/examples/brain-functional-networks/).
+
+## Documentation
+
+| Page | Content |
+| --- | --- |
+| [Getting started](https://dubris.github.io/pavi-pyro/getting-started/) | installation and a first model |
+| [Concepts](https://dubris.github.io/pavi-pyro/concepts/) | hierarchical models, variational inference, normalizing flows, plate amortization |
+| [The PAVI guide](https://dubris.github.io/pavi-pyro/guide/) | the arguments of `PAVIGuide` |
+| [Encodings](https://dubris.github.io/pavi-pyro/encodings/) · [Transforms](https://dubris.github.io/pavi-pyro/transforms/) · [Schemes](https://dubris.github.io/pavi-pyro/schemes/) | the building blocks and their options |
+| [Training](https://dubris.github.io/pavi-pyro/training/) | ELBOs, enumeration, mini-batches, Lightning, KL annealing |
+| [Examples](https://dubris.github.io/pavi-pyro/examples/) | toy model, brain functional networks |
+
+The documentation is published at <https://dubris.github.io/pavi-pyro/>. To browse it locally, run `uv run --group docs mkdocs serve`.
+
+## Citation
+
+If you use `pavi`, please cite:
+
+```bibtex
+@article{rouillard2023pavi,
+  title   = {{PAVI}: Plate-Amortized Variational Inference},
+  author  = {Rouillard, Louis and Le Bris, Alexandre and Moreau, Thomas and Wassermann, Demian},
+  journal = {Transactions on Machine Learning Research},
+  year    = {2023}
+}
+```
+
+For the brain functional network application:
+
+```bibtex
+@unpublished{lebris2024transfer,
+  title  = {Improving Individual-Specific Functional Parcellation through Transfer Learning},
+  author = {Le Bris, Alexandre and Rouillard, Louis and Wassermann, Demian},
+  year   = {2024},
+  url    = {https://hal.science/hal-04506561}
+}
+```
+
+## License
+
+[MIT](https://github.com/Dubris/pavi-pyro/blob/main/LICENSE)
+
+[^plate]: **Plate**: a set of random variables that are conditionally independent given their parents, e.g. one variable per subject or per brain vertex. Declared with `pyro.plate`.
+[^guide]: **Guide**: in Pyro, the parameterized family of distributions $q$ (the *variational family*) that is fitted to approximate the posterior. `pavi` provides the guide; you provide the model.
+[^site]: **Site**: a random variable of a Pyro model, declared with `pyro.sample("name", distribution)`. Latent sites are inferred; observed sites hold the data.

pavi_pyro-0.1.0/README.md

@@ -0,0 +1,196 @@
+# pavi — plate-amortized variational inference in Pyro
+
+**Scalable Bayesian inference for hierarchical models with millions of latent variables.**
+
+[![Documentation](https://img.shields.io/badge/docs-GitHub%20Pages-blue)](https://dubris.github.io/pavi-pyro/)
+[![Python](https://img.shields.io/badge/python-%E2%89%A53.14-3776ab)](https://github.com/Dubris/pavi-pyro/blob/main/pyproject.toml)
+[![Pyro](https://img.shields.io/badge/built%20on-Pyro%20%2B%20PyTorch-ee4c2c)](https://pyro.ai)
+[![License: MIT](https://img.shields.io/badge/license-MIT-green)](https://github.com/Dubris/pavi-pyro/blob/main/LICENSE)
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/parcellations_group_vs_individuals.jpg" width="85%" alt="Population-level and individual parcellations of the cortex into 17 functional networks, estimated with PAVI on the HCP dataset">
+  <br>
+  <em>Population-level (top) and individual (bottom) parcellations of the cortex into 17 functional networks, inferred with PAVI from HCP resting-state fMRI (<a href="https://hal.science/hal-04506561">Le Bris et al., 2024</a>).</em>
+</p>
+
+`pavi` implements **PAVI** ([Rouillard et al., 2023](#citation)) for [Pyro](https://pyro.ai), the probabilistic programming library built on PyTorch. You describe your data with a hierarchical Bayesian model; `pavi` approximates the posterior distribution of its latent variables, even when there are millions of them. In particular, a main application of the algorithm and this implementation was estimating individualized functional networks over the cortex from fMRI data (see the figure above and [Le Bris et al., 2024](https://hal.science/hal-04506561)).
+
+## Why PAVI?
+
+Hierarchical models repeat the same variable across plates[^plate], and the number of latent variables quickly explodes. In the brain application, about 1,000 subjects × 60,000 cortical vertices give **60 million** latent variables. With a classical mean-field Gaussian approximation over about 20 functional networks, this requires **more than 2 billion** variational parameters.
+
+PAVI builds the variational family with **normalizing flows**: a simple base distribution is mapped to a complex one by an invertible, learnable transform $f$. For an instance $n$ of a plate,
+
+$$z^n = f(u^n;\ \phi, E^n), \qquad u^n \sim \text{base distribution}.$$
+
+- The parameters $\phi$ of the transform are **shared** by all the instances of the plate. They capture what instances have in common and act as a group-level regularization.
+- Each instance only keeps a small **encoding** $E^n \in \mathbb{R}^m$, which encodes the statistics of its posterior (hence the name).
+
+The number of parameters still grows with the number of instances $N$, but much more slowly: for large $N$ it is reduced by a factor $O(m/d)$ compared to the mean-field approximation, where $d$ is the size of each variable and $m \ll d$. The mean-field approximation itself is the special case of a standard normal base and a diagonal affine transform. New subjects can be inferred by **refitting only their encodings**.
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/pavi_plate_amortization.png" width="45%" alt="(A) Generative model and PAVI variational family: posteriors of z_1 are parameterized by a shared phi_1 and per-instance encodings E_1^n. (B) Each posterior is a base distribution mapped by an invertible transform f.">
+</p>
+
+Learn more in [Concepts](https://dubris.github.io/pavi-pyro/concepts/).
+
+## Features
+
+- **Bring your own generative model**, implemented with Pyro and PyTorch.
+- **An off-the-shelf implementation of PAVI**, with training utilities: plate-aware mini-batch sampler, PyTorch Lightning inference module, KL annealing and logging callbacks.
+- **Many customization options**:
+  - a wide choice of parameterized invertible [transforms](https://dubris.github.io/pavi-pyro/transforms/);
+  - several strategies to manage the [encodings](https://dubris.github.io/pavi-pyro/encodings/);
+  - several strategies for the base distribution, called [schemes](https://dubris.github.io/pavi-pyro/schemes/).
+
+## Installation
+
+`pavi` requires Python ≥ 3.14. It is published on PyPI as `pavi-pyro` and imported as `pavi`:
+
+```bash
+pip install pavi-pyro  # or: uv add pavi-pyro
+```
+
+See [Getting started](https://dubris.github.io/pavi-pyro/getting-started/) to install from source.
+
+## Quickstart
+
+A three-level hierarchical model: a global mean $z_2$, the means $z_1^i$ of 5 groups (plate `P_1`), and 50 observations $X^{i,j}$ per group (plate `P_0`).
+
+**1. Write the generative model in Pyro.** The model receives the data and, optionally, the indices of a mini-batch (`idx`).
+
+```python
+import pyro
+import pyro.distributions as dist
+import torch
+
+from pavi.core import JointEM, NormalizingFlowScheme, PAVIGuide
+from pavi.distributions import affine, conditional_affine
+
+NUM_GROUPS, NUM_OBS, DIM = 5, 50, 2
+
+
+def model(data=None, idx=None):
+    idx = idx or {}
+    z_2 = pyro.sample("z_2", dist.Normal(torch.zeros(DIM), 1.0).to_event(1))
+    with pyro.plate("P_1", NUM_GROUPS, subsample=idx.get("P_1"), dim=-1):
+        z_1 = pyro.sample("z_1", dist.Normal(z_2, 5.0).to_event(1))
+        with pyro.plate("P_0", NUM_OBS, subsample=idx.get("P_0"), dim=-2):
+            pyro.sample("X", dist.Normal(z_1, 1.0).to_event(1), obs=data)
+
+
+# Simulate a dataset from the model
+pyro.set_rng_seed(0)
+X = pyro.poutine.trace(model).get_trace().nodes["X"]["value"]
+```
+
+**2. Build the PAVI guide[^guide].** Choose a transform for each latent site[^site], an encoding manager and a scheme.
+
+```python
+guide = PAVIGuide(
+    model,
+    transform_modules={
+        "z_2": [affine(scale_init=1.0)],                     # global variable
+        "z_1": [conditional_affine(hidden_layers=[32, 32])],  # plate P_1: conditioned on encodings
+    },
+    transform_scheme=NormalizingFlowScheme,
+    encoding_manager=JointEM(encoding_size={"P_1": 8}),
+    hidden_sites=["X"],  # observed sites are not inferred
+    device="cpu",
+    data=X,
+)
+```
+
+**3. Fit the guide** by maximizing the evidence lower bound (ELBO).
+
+```python
+elbo = pyro.infer.Trace_ELBO(num_particles=8)
+optimizer = torch.optim.Adam(guide.parameters(), lr=1e-2)
+for step in range(2000):
+    optimizer.zero_grad()
+    loss = elbo.differentiable_loss(model, guide, X)
+    loss.backward()
+    optimizer.step()
+```
+
+**4. Sample from the approximate posterior.**
+
+```python
+with torch.no_grad():
+    posterior_sample = guide(X)
+print(posterior_sample["z_1"])  # one sample of the 5 group means
+```
+
+The [toy-model notebook](https://dubris.github.io/pavi-pyro/examples/toymodel_hgm/) compares the result with the exact posterior.
+
+## How it works
+
+```mermaid
+flowchart LR
+    M["Pyro model<br/>with plates"] -->|"inspected once"| S["latent sites<br/>and their plates"]
+    S --> EM["encoding manager<br/>encodings E"]
+    S --> T["transforms<br/>per site"]
+    EM -->|"conditions"| T
+    T --> SC["scheme<br/>base distribution"]
+    SC --> Q["approximate posterior q"]
+```
+
+When it is created, the guide runs the model once to list its latent sites and their plates. The encoding manager then creates the encodings, and each site gets its chain of transforms, conditioned on the encodings when the site lives in a plate. The scheme chooses the base distribution that the transforms are applied to. See [The PAVI guide](https://dubris.github.io/pavi-pyro/guide/).
+
+## Application: individual brain functional networks
+
+PAVI makes it possible to map the **functional networks of each individual brain** at the resolution of the cortical mesh, from resting-state fMRI ([Le Bris et al., 2024](https://hal.science/hal-04506561)). A hierarchical model explains the functional connectivity of each vertex as a mixture of 17 networks, whose spatial organization varies across individuals around a population-level atlas.
+
+Instead of a discrete atlas, the model yields a **continuous, probabilistic map of each network**, at the population level (below) and for each individual:
+
+<p align="center">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/probabilistic_networks_1.jpg" width="48%" alt="Population-level probabilistic maps of functional networks VIN-1 to CON-2">
+  <img src="https://raw.githubusercontent.com/Dubris/pavi-pyro/main/docs/figures/probabilistic_networks_2.jpg" width="48%" alt="Population-level probabilistic maps of functional networks CON-3 to DMN-4">
+</p>
+
+These individual maps are reliable, reflect individual-specific organization, predict behavior, and transfer to small datasets through the population-level priors. See [the brain application example](https://dubris.github.io/pavi-pyro/examples/brain-functional-networks/).
+
+## Documentation
+
+| Page | Content |
+| --- | --- |
+| [Getting started](https://dubris.github.io/pavi-pyro/getting-started/) | installation and a first model |
+| [Concepts](https://dubris.github.io/pavi-pyro/concepts/) | hierarchical models, variational inference, normalizing flows, plate amortization |
+| [The PAVI guide](https://dubris.github.io/pavi-pyro/guide/) | the arguments of `PAVIGuide` |
+| [Encodings](https://dubris.github.io/pavi-pyro/encodings/) · [Transforms](https://dubris.github.io/pavi-pyro/transforms/) · [Schemes](https://dubris.github.io/pavi-pyro/schemes/) | the building blocks and their options |
+| [Training](https://dubris.github.io/pavi-pyro/training/) | ELBOs, enumeration, mini-batches, Lightning, KL annealing |
+| [Examples](https://dubris.github.io/pavi-pyro/examples/) | toy model, brain functional networks |
+
+The documentation is published at <https://dubris.github.io/pavi-pyro/>. To browse it locally, run `uv run --group docs mkdocs serve`.
+
+## Citation
+
+If you use `pavi`, please cite:
+
+```bibtex
+@article{rouillard2023pavi,
+  title   = {{PAVI}: Plate-Amortized Variational Inference},
+  author  = {Rouillard, Louis and Le Bris, Alexandre and Moreau, Thomas and Wassermann, Demian},
+  journal = {Transactions on Machine Learning Research},
+  year    = {2023}
+}
+```
+
+For the brain functional network application:
+
+```bibtex
+@unpublished{lebris2024transfer,
+  title  = {Improving Individual-Specific Functional Parcellation through Transfer Learning},
+  author = {Le Bris, Alexandre and Rouillard, Louis and Wassermann, Demian},
+  year   = {2024},
+  url    = {https://hal.science/hal-04506561}
+}
+```
+
+## License
+
+[MIT](https://github.com/Dubris/pavi-pyro/blob/main/LICENSE)
+
+[^plate]: **Plate**: a set of random variables that are conditionally independent given their parents, e.g. one variable per subject or per brain vertex. Declared with `pyro.plate`.
+[^guide]: **Guide**: in Pyro, the parameterized family of distributions $q$ (the *variational family*) that is fitted to approximate the posterior. `pavi` provides the guide; you provide the model.
+[^site]: **Site**: a random variable of a Pyro model, declared with `pyro.sample("name", distribution)`. Latent sites are inferred; observed sites hold the data.

pavi_pyro-0.1.0/pyproject.toml

@@ -0,0 +1,137 @@
+[project]
+name = "pavi-pyro"
+version = "0.1.0"
+description = "Plate-amortized variational inference (PAVI) for large hierarchical Bayesian models, built on Pyro"
+readme = "README.md"
+license = "MIT"
+license-files = ["LICENSE"]
+keywords = [
+    "variational inference",
+    "hierarchical Bayesian models",
+    "pyro",
+    "normalizing flows",
+    "neuroimaging",
+]
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Operating System :: OS Independent",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.14",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+    "Topic :: Scientific/Engineering :: Mathematics",
+    "Typing :: Typed",
+]
+requires-python = ">=3.14"
+dependencies = [
+    "numpy>=2.5.3",
+    "pyro-ppl>=1.9.1",
+    "pytorch-lightning>=2.6.5",
+    "torch>=2.15.0",
+    "zuko>=1.6.0",
+]
+
+[[project.authors]]
+name = "Alexandre Le Bris"
+email = "alexandrelebris@gmail.com"
+
+[project.urls]
+Homepage = "https://github.com/Dubris/pavi-pyro"
+Documentation = "https://dubris.github.io/pavi-pyro/"
+Repository = "https://github.com/Dubris/pavi-pyro"
+Issues = "https://github.com/Dubris/pavi-pyro/issues"
+
+[build-system]
+requires = ["uv_build>=0.11.25,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+module-name = "pavi"
+
+[tool.ruff]
+line-length = 140
+include = [
+    "*.py",
+    "*.pyi",
+    "**/__init__.py",
+    "*.ipynb",
+]
+exclude = [
+    ".git",
+    ".venv",
+    "__pycache__",
+    "build",
+    "dist",
+]
+extend-exclude = [
+    "*.interactive",
+    "untitled:*",
+]
+
+[tool.ruff.lint]
+select = [
+    "E",
+    "F",
+    "W",
+    "I",
+    "B",
+    "UP",
+    "NPY",
+    "C4",
+]
+ignore = [
+    "E501",
+    "B008",
+]
+extend-select = [
+    "TD",
+    "FIX",
+]
+
+[tool.vscode.settings]
+"jupyter.notebookFileRoot" = "${workspaceFolder}"
+
+[tool.mypy]
+files = ["src"]
+python_version = "3.14"
+
+[[tool.mypy.overrides]]
+module = ["zuko.*"]
+ignore_missing_imports = true
+
+[tool.pytest.ini_options]
+addopts = "-s"
+testpaths = ["test"]
+
+[dependency-groups]
+dev = [
+    "mypy>=2.4.1",
+    "pytest>=9.1.1",
+    "pytest-mock>=3.15.1",
+]
+docs = [
+    "markdown-callouts>=0.4.0",
+    "mkdocs-jupyter>=0.27.0",
+    "mkdocs-material>=9.8.0",
+    "mkdocstrings[python]>=1.0.8",
+]
+experiments = [
+    "hcp-utils>=0.1.0",
+    "psutil>=7.2.2",
+    "hydra-core==1.3.5",
+    "ipykernel>=7.3.0",
+    "nibabel>=5.4.2",
+    "nilearn>=0.13.1",
+    "omegaconf>=2.3.1",
+    "seaborn>=0.13.2",
+    "xarray>=2026.7.0",
+    "lightning>=2.6.5",
+    "mlflow>=3.15.2",
+    "hydra-zen>=0.16.0",
+    "beartype>=0.22.9",
+    "einops>=0.8.2",
+    "brainspace>=0.2.1",
+    "typer>=0.27.2",
+    "graphviz>=0.21",
+]