From f355fa5440fafd11290b6e31542838bd56cfab47 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sun, 9 Aug 2026 11:45:12 +0530 Subject: [PATCH 01/29] list the estimator configs in the API reference --- docs/api_reference/neural_nets.rst | 100 +++++++++++++++++++++++++++++ 1 file changed, 100 insertions(+) diff --git a/docs/api_reference/neural_nets.rst b/docs/api_reference/neural_nets.rst index d3d72698e..11089af5d 100644 --- a/docs/api_reference/neural_nets.rst +++ b/docs/api_reference/neural_nets.rst @@ -4,6 +4,106 @@ Neural nets =========== +The neural network of an inference method is described by a *config* object: one +class per model, holding only the settings that model accepts. Pass it to the +trainer as ``density_estimator``, ``classifier``, or ``vf_estimator``. See +:doc:`../how_to_guide/27_estimator_configs` for how to use them. + + +Density estimator configs +------------------------- + +For ``NPE`` and ``NLE``, passed as ``density_estimator``. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.DensityConfigBase + sbi.neural_nets.MADEConfig + sbi.neural_nets.MAFConfig + sbi.neural_nets.MAFRQSConfig + sbi.neural_nets.MDNConfig + sbi.neural_nets.NSFConfig + sbi.neural_nets.TabPFNConfig + sbi.neural_nets.ZukoBPFConfig + sbi.neural_nets.ZukoGFConfig + sbi.neural_nets.ZukoMAFConfig + sbi.neural_nets.ZukoNAFConfig + sbi.neural_nets.ZukoNCSFConfig + sbi.neural_nets.ZukoNICEConfig + sbi.neural_nets.ZukoNSFConfig + sbi.neural_nets.ZukoSOSPFConfig + sbi.neural_nets.ZukoUNAFConfig + + +Mixed density estimator configs +------------------------------- + +For ``MNPE`` and ``MNLE``, where part of the data is discrete. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.MixedConfig + + +Classifier configs +------------------ + +For the ``NRE`` variants, passed as ``classifier``. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.ClassifierConfigBase + sbi.neural_nets.LinearClassifierConfig + sbi.neural_nets.MLPClassifierConfig + sbi.neural_nets.ResNetClassifierConfig + + +Marginal density estimator configs +---------------------------------- + +For ``MarginalTrainer``, which fits an unconditional density. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.MarginalConfigBase + sbi.neural_nets.MarginalBPFConfig + sbi.neural_nets.MarginalGFConfig + sbi.neural_nets.MarginalMAFConfig + sbi.neural_nets.MarginalNAFConfig + sbi.neural_nets.MarginalNCSFConfig + sbi.neural_nets.MarginalNICEConfig + sbi.neural_nets.MarginalNSFConfig + sbi.neural_nets.MarginalSOSPFConfig + sbi.neural_nets.MarginalUNAFConfig + + +Vector field estimator builders +------------------------------- + +For ``FMPE`` and ``NPSE``, passed as ``vf_estimator``. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.VectorFieldEstimatorBuilder + + +Factory functions +----------------- + +The factory functions predate the config classes and are kept for backwards +compatibility. They take the model name as a string and return a build +function. + .. autosummary:: :toctree: _autosummary :nosignatures: From 7439de29dd95716d16dcb20f087989fd103d5638 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sun, 9 Aug 2026 14:12:33 +0530 Subject: [PATCH 02/29] add a how-to for the estimator configs --- docs/how_to_guide/27_estimator_configs.ipynb | 497 +++++++++++++++++++ 1 file changed, 497 insertions(+) create mode 100644 docs/how_to_guide/27_estimator_configs.ipynb diff --git a/docs/how_to_guide/27_estimator_configs.ipynb b/docs/how_to_guide/27_estimator_configs.ipynb new file mode 100644 index 000000000..b87be414d --- /dev/null +++ b/docs/how_to_guide/27_estimator_configs.ipynb @@ -0,0 +1,497 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# How to configure the neural network" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Every inference method in `sbi` trains a neural network. Which network it trains, and with which settings, is described by a **config object** that you pass to the trainer.\n", + "\n", + "There is one config class per model, and it carries only the settings that model actually accepts. A setting the model does not have is therefore not a field on its config, so passing it raises an error right away instead of being silently ignored.\n", + "\n", + "This guide covers:\n", + "\n", + "- picking a model and changing its hyperparameters\n", + "- adding an embedding network and controlling z-scoring\n", + "- the configs for classifiers, mixed data, and marginal densities\n", + "- moving off the deprecated string and factory-function interfaces" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "import torch\n", + "\n", + "from sbi.utils import BoxUniform\n", + "\n", + "num_dim = 2\n", + "prior = BoxUniform(low=-2 * torch.ones(num_dim), high=2 * torch.ones(num_dim))\n", + "\n", + "theta = prior.sample((500,))\n", + "x = theta + 1.0 + 0.1 * torch.randn_like(theta)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Choosing a density estimator\n", + "\n", + "`NPE` and `NLE` take their network as `density_estimator`. Import the config of the model you want and pass an instance:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.inference import NPE\n", + "from sbi.neural_nets import ZukoNSFConfig\n", + "\n", + "trainer = NPE(prior=prior, density_estimator=ZukoNSFConfig())" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Hyperparameters are constructor arguments, so your editor can complete them and type checkers can see them:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ZukoNSFConfig(hidden_features=64, num_transforms=8, num_bins=12)\n", + "trainer = NPE(prior=prior, density_estimator=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The available density estimator configs are:\n", + "\n", + "| Config | Model |\n", + "|---|---|\n", + "| `MAFConfig` | masked autoregressive flow (`nflows`) |\n", + "| `MAFRQSConfig` | MAF with rational-quadratic splines (`nflows`) |\n", + "| `NSFConfig` | neural spline flow (`nflows`) |\n", + "| `MADEConfig` | masked autoencoder for density estimation |\n", + "| `MDNConfig` | mixture density network |\n", + "| `TabPFNConfig` | TabPFN-based estimator |\n", + "| `ZukoMAFConfig`, `ZukoNSFConfig`, `ZukoNCSFConfig`, `ZukoNAFConfig`, `ZukoUNAFConfig`, `ZukoBPFConfig`, `ZukoSOSPFConfig`, `ZukoNICEConfig`, `ZukoGFConfig` | the corresponding [`zuko`](https://github.com/probabilists/zuko) flows |\n", + "\n", + "The `nflows` package is no longer maintained, so for new projects we tentatively recommend the `zuko` variants." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "`NLE` estimates the likelihood rather than the posterior, but it takes the same configs, because the difference is which variable is modeled and which is conditioned on. That is decided by the trainer, not by the config:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.inference import NLE\n", + "from sbi.neural_nets import MAFConfig\n", + "\n", + "trainer = NLE(prior=prior, density_estimator=MAFConfig(hidden_features=64))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Settings a model does not have\n", + "\n", + "`num_bins` is a spline setting. `NSFConfig` has it, `MAFConfig` does not, so the mistake surfaces immediately:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.neural_nets import NSFConfig\n", + "\n", + "NSFConfig(num_bins=12) # fine, a spline flow has bins\n", + "\n", + "try:\n", + " MAFConfig(num_bins=12)\n", + "except TypeError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The same applies to a misspelled setting, and to a value outside the allowed set:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " NSFConfig(hiden_features=64)\n", + "except TypeError as e:\n", + " print(e)\n", + "\n", + "try:\n", + " NSFConfig(z_score_input=\"strucured\")\n", + "except ValueError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Inspecting a config\n", + "\n", + "Configs are frozen dataclasses, and their `repr` shows only what you changed, which makes them convenient to log or to put in an experiment record:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "print(ZukoNSFConfig())\n", + "print(ZukoNSFConfig(hidden_features=64, num_transforms=8))" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Being frozen means a config cannot be edited after construction. To vary one setting, build a new config:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "configs = [ZukoNSFConfig(hidden_features=h) for h in (32, 64, 128)]\n", + "print(configs)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Embedding networks\n", + "\n", + "For high-dimensional data, pass an embedding network that learns summary statistics of the conditioning variable. See [how to use embedding networks](04_embedding_networks.ipynb) for how to choose one." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.neural_nets.embedding_nets import FCEmbedding\n", + "\n", + "config = ZukoNSFConfig(embedding_net=FCEmbedding(input_dim=num_dim, output_dim=8))\n", + "trainer = NPE(prior=prior, density_estimator=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Z-scoring\n", + "\n", + "Both variables are z-scored independently by default. The standardization is built into the estimator, so you keep passing raw data at inference time.\n", + "\n", + "`z_score_input` applies to the modeled variable and `z_score_condition` to the variable conditioned on. Which is which depends on the method: for `NPE` the input is $\\theta$ and the condition is $x$, for `NLE` it is the other way round. Each takes `\"independent\"` (default), `\"structured\"`, or `\"none\"`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ZukoNSFConfig(z_score_input=\"independent\", z_score_condition=\"structured\")\n", + "trainer = NPE(prior=prior, density_estimator=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Use `\"structured\"` when the entries of a variable are not exchangeable, for example a time series, where a single mean and standard deviation across all entries would be the wrong summary. Use `\"none\"` if the data is already standardized." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Settings without a field\n", + "\n", + "Some underlying models accept keyword arguments that have no field on the config. `extra_kwargs` forwards them, and marks at the call site that you are stepping outside the checked surface:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ZukoNSFConfig(hidden_features=64, extra_kwargs={\"randperm\": True})\n", + "trainer = NPE(prior=prior, density_estimator=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "A key that duplicates an existing field is rejected, so there is one place a setting can come from:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "try:\n", + " ZukoNSFConfig(extra_kwargs={\"hidden_features\": 64})\n", + "except ValueError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Classifiers for NRE\n", + "\n", + "The `NRE` variants train a classifier instead of a density estimator, and take it as `classifier`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.inference import NRE\n", + "from sbi.neural_nets import ResNetClassifierConfig\n", + "\n", + "config = ResNetClassifierConfig(hidden_features=64, num_blocks=3)\n", + "trainer = NRE(prior=prior, classifier=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The classifier configs are `LinearClassifierConfig`, `MLPClassifierConfig`, and `ResNetClassifierConfig`. A classifier sees both variables as inputs rather than conditioning on one, so it takes two embedding networks instead of one:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "config = ResNetClassifierConfig(\n", + " embedding_net_theta=FCEmbedding(input_dim=num_dim, output_dim=8),\n", + " embedding_net_x=FCEmbedding(input_dim=num_dim, output_dim=8),\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Mixed data for MNPE and MNLE\n", + "\n", + "`MNPE` and `MNLE` handle data that is partly continuous and partly discrete. `MixedConfig` describes the whole estimator, and the continuous part is configured by nesting that model's own config, so its settings stay validated by its own class:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.neural_nets import MixedConfig\n", + "\n", + "config = MixedConfig(\n", + " continuous=ZukoNSFConfig(hidden_features=64, num_transforms=4),\n", + " discrete_hidden_features=32,\n", + ")\n", + "print(config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "Note that z-scoring of the modeled variable is a property of the continuous model, so it lives on the nested config, while `z_score_condition` stays on `MixedConfig`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Marginal densities\n", + "\n", + "`MarginalTrainer` fits an unconditional density, so its configs have no conditioning variable and no `z_score_condition`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.inference.trainers.marginal import MarginalTrainer\n", + "from sbi.neural_nets import MarginalNSFConfig\n", + "\n", + "trainer = MarginalTrainer(\n", + " density_estimator=MarginalNSFConfig(hidden_features=64, num_transforms=4)\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Building the estimator yourself\n", + "\n", + "The trainer calls `build` on the config once it has seen data, because the network needs the shapes and the z-scoring statistics. You can call it yourself if you want the network outside a trainer, for example for a custom training loop as in [the training interface tutorial](../advanced_tutorials/18_training_interface.ipynb):" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "estimator = ZukoNSFConfig(hidden_features=32).build(theta, x)\n", + "print(type(estimator).__name__, \"| input shape\", estimator.input_shape)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The first argument is the modeled variable and the second the one conditioned on, matching `z_score_input` and `z_score_condition`." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Moving off strings and factory functions\n", + "\n", + "Before the configs, the network was selected with a string, or with one of the `posterior_nn` / `likelihood_nn` / `classifier_nn` factory functions. Both still work; passing a string now warns.\n", + "\n", + "| Before | Now |\n", + "|---|---|\n", + "| `NPE(prior, density_estimator=\"nsf\")` | `NPE(prior, density_estimator=NSFConfig())` |\n", + "| `NPE(prior, density_estimator=\"zuko_nsf\")` | `NPE(prior, density_estimator=ZukoNSFConfig())` |\n", + "| `NLE(prior, density_estimator=\"maf\")` | `NLE(prior, density_estimator=MAFConfig())` |\n", + "| `NRE(prior, classifier=\"resnet\")` | `NRE(prior, classifier=ResNetClassifierConfig())` |\n", + "| `posterior_nn(model=\"nsf\", hidden_features=64)` | `NSFConfig(hidden_features=64)` |\n", + "| `likelihood_nn(model=\"maf\", num_transforms=8)` | `MAFConfig(num_transforms=8)` |\n", + "| `classifier_nn(model=\"mlp\", hidden_features=64)` | `MLPClassifierConfig(hidden_features=64)` |\n", + "\n", + "The model name moves from a string argument into the class name, and the remaining keyword arguments carry over unchanged." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.neural_nets import posterior_nn\n", + "\n", + "# Before\n", + "build_fn = posterior_nn(model=\"zuko_nsf\", hidden_features=64, num_transforms=8)\n", + "trainer = NPE(prior=prior, density_estimator=build_fn)\n", + "\n", + "# Now\n", + "config = ZukoNSFConfig(hidden_features=64, num_transforms=8)\n", + "trainer = NPE(prior=prior, density_estimator=config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "One difference worth knowing when you migrate: the factory functions accept settings the chosen model ignores, whereas the config rejects them. If a call that used to run now raises a `TypeError`, that setting was not reaching the network before either." + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## See also\n", + "\n", + "- [How to choose neural nets](03_choose_neural_net.ipynb) for which model to pick\n", + "- [How to use embedding networks](04_embedding_networks.ipynb)\n", + "- [How to choose abstraction levels](24_abstraction_levels.ipynb) for where configs sit relative to the other ways of specifying a network" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.4" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} From 1b8e76866b2ee882f032c74f02a60fa82cdaddca Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sun, 9 Aug 2026 16:22:11 +0530 Subject: [PATCH 03/29] link the config how-to from the neural nets section --- docs/how_to_guide.rst | 1 + docs/how_to_guide/neural_nets.rst | 1 + 2 files changed, 2 insertions(+) diff --git a/docs/how_to_guide.rst b/docs/how_to_guide.rst index cb2afc2c2..ace600424 100644 --- a/docs/how_to_guide.rst +++ b/docs/how_to_guide.rst @@ -34,6 +34,7 @@ Neural nets ----------- - :doc:`how_to_guide/03_choose_neural_net` +- :doc:`how_to_guide/27_estimator_configs` - :doc:`how_to_guide/04_embedding_networks` - :doc:`how_to_guide/20_time_series_embedding` - :doc:`how_to_guide/08_permutation_invariant_embeddings` diff --git a/docs/how_to_guide/neural_nets.rst b/docs/how_to_guide/neural_nets.rst index 75407d522..5b90005ed 100644 --- a/docs/how_to_guide/neural_nets.rst +++ b/docs/how_to_guide/neural_nets.rst @@ -8,6 +8,7 @@ Neural nets :maxdepth: 1 03_choose_neural_net.ipynb + 27_estimator_configs.ipynb 04_embedding_networks.ipynb 20_time_series_embedding.ipynb 08_permutation_invariant_embeddings.ipynb From 5437eaed4be09eeeb8ec91baecc5a6dc9ad4ec79 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Mon, 10 Aug 2026 11:15:00 +0530 Subject: [PATCH 04/29] rework the abstraction levels around the config objects --- docs/how_to_guide/24_abstraction_levels.ipynb | 68 ++++++++++--------- 1 file changed, 37 insertions(+), 31 deletions(-) diff --git a/docs/how_to_guide/24_abstraction_levels.ipynb b/docs/how_to_guide/24_abstraction_levels.ipynb index d9c629d2d..350e4d344 100644 --- a/docs/how_to_guide/24_abstraction_levels.ipynb +++ b/docs/how_to_guide/24_abstraction_levels.ipynb @@ -75,11 +75,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Level 1: Trainer Classes (Recommended)\n", + "## Level 1: Trainer Defaults (Recommended Start)\n", "\n", "**Use case**: Standard workflows, most common approach\n", "\n", - "The trainer classes provide the recommended interface with string-based customization." + "Construct a trainer and train. Each method ships a default network that works reasonably across problems, so nothing has to be specified." ] }, { @@ -88,8 +88,8 @@ "metadata": {}, "outputs": [], "source": [ - "# Level 1: Simple trainer class with string specification\n", - "inference = NPE(prior=prior, density_estimator=\"nsf\")\n", + "# Level 1: trainer defaults\n", + "inference = NPE(prior=prior)\n", "\n", "# Train on the data\n", "inference.append_simulations(theta, x)\n", @@ -99,7 +99,7 @@ "posterior = inference.build_posterior()\n", "samples_lvl1 = posterior.sample((1000,), x=x_o)\n", "\n", - "print(\"Level 1 complete - used NSF with default settings\")" + "print(\"Level 1 complete - used the default density estimator\")" ] }, { @@ -107,7 +107,7 @@ "metadata": {}, "source": [ "**Key features**:\n", - "- Simple string specification: `\"nsf\"`, `\"maf\"`, `\"zuko_nsf\"`, `\"mdn\"`, etc.\n", + "- No network configuration at all\n", "- Multi-round inference support\n", "- Automatic handling of training loops\n", "- **Start here** for most use cases" @@ -117,11 +117,11 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "## Level 2: Factory Functions\n", + "## Level 2: Config Objects\n", "\n", - "**Use case**: Need specific architecture hyperparameters\n", + "**Use case**: Choose the model, or tune its architecture hyperparameters\n", "\n", - "Use factory functions like `posterior_nn()` when you need to tune the network architecture." + "Pass a config object as `density_estimator`. There is one class per model, and it carries exactly the settings that model accepts, so an unsupported or misspelled setting is an error at construction rather than something silently dropped." ] }, { @@ -130,13 +130,12 @@ "metadata": {}, "outputs": [], "source": [ - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import MAFConfig\n", "\n", - "# Level 2: Factory function with custom hyperparameters\n", - "density_estimator = posterior_nn(\n", - " model=\"maf\", # Masked Autoregressive Flow\n", - " hidden_features=50, # Customize hidden layer size\n", - " num_transforms=5, # Customize number of transform layers\n", + "# Level 2: config object with custom hyperparameters\n", + "density_estimator = MAFConfig(\n", + " hidden_features=50, # Customize hidden layer size\n", + " num_transforms=5, # Customize number of transform layers\n", ")\n", "\n", "# Pass to NPE (rest of workflow is the same)\n", @@ -155,10 +154,14 @@ "metadata": {}, "source": [ "**Key features**:\n", - "- Fine-grained control over hyperparameters\n", + "- Fine-grained control over hyperparameters, checked at construction\n", "- Can add embedding networks for high-dimensional data\n", "- Still benefits from trainer conveniences\n", - "- For NPE: `posterior_nn()`, for NLE: `likelihood_nn()`, for NRE: `classifier_nn()`" + "- Same configs for NPE and NLE; NRE takes classifier configs as `classifier`\n", + "\n", + "See [how to configure the neural network](27_estimator_configs.ipynb) for the full list of configs and their settings.\n", + "\n", + "The older interfaces for this level still work: a model name as a string, or the `posterior_nn()` / `likelihood_nn()` / `classifier_nn()` factory functions. Strings are deprecated and warn; the migration table is in the same guide." ] }, { @@ -167,9 +170,9 @@ "source": [ "## Level 3: Direct Network Builders\n", "\n", - "**Use case**: Custom neural network architecture with full parameter access\n", + "**Use case**: Reach a build-function argument that no config exposes\n", "\n", - "Use direct builder functions like `build_nsf()` for maximum control over network construction." + "Pass any callable taking the two training batches. The `build_*` functions are such callables once their settings are bound, which gives access to their full signature." ] }, { @@ -182,7 +185,7 @@ "\n", "from sbi.neural_nets.net_builders.flow import build_nsf\n", "\n", - "# Level 3: Direct builder with full parameter control\n", + "# Level 3: bind the build function's own arguments\n", "custom_builder = partial(\n", " build_nsf,\n", " hidden_features=60,\n", @@ -209,7 +212,9 @@ "**Key features**:\n", "- Direct access to all builder parameters\n", "- Maximum flexibility for architecture design\n", - "- Can implement fully custom architectures by subclassing `DensityEstimator`" + "- Can implement fully custom architectures by subclassing `DensityEstimator`\n", + "\n", + "For a setting the config does not have but the build function does, `extra_kwargs` on the config is usually the smaller step: `ZukoMAFConfig(extra_kwargs={\"randperm\": True})`." ] }, { @@ -247,7 +252,7 @@ "- Must combine with prior using MCMC, VI, or rejection sampling to get posterior samples\n", "- This adds a **second dimension of control**: choosing and configuring the sampling method\n", "\n", - "**Important**: The 4 density estimator levels above work exactly the same for NLE - just use `likelihood_nn()` instead of `posterior_nn()` at Level 2." + "**Important**: The 4 density estimator levels above work exactly the same for NLE, with the same config classes at Level 2. What changes is which variable is modeled and which is conditioned on, and the trainer decides that, not the config." ] }, { @@ -266,7 +271,7 @@ "outputs": [], "source": [ "# Level 1 with NLE - same pattern as NPE\n", - "inference_nle = NLE(prior=prior, density_estimator=\"nsf\")\n", + "inference_nle = NLE(prior=prior)\n", "inference_nle.append_simulations(theta, x)\n", "likelihood_net = inference_nle.train()\n", "\n", @@ -283,7 +288,7 @@ "metadata": {}, "source": [ "**Note**: Levels 2-4 for the density estimator work identically:\n", - "- Level 2: Use `likelihood_nn()` instead of `posterior_nn()`\n", + "- Level 2: the same config classes, e.g. `NLE(prior, density_estimator=MAFConfig())`\n", "- Level 3: Use `build_nsf()` (same as NPE)\n", "- Level 4: Custom training (see tutorial 18)" ] @@ -507,11 +512,12 @@ "\n", "| You want to... | Use Level | Example |\n", "|----------------|-----------|----------|\n", - "| Standard workflows with good defaults | **1** | `NPE(prior, density_estimator=\"nsf\")` |\n", - "| Try different density estimator types | **1** | Switch `\"nsf\"`, `\"maf\"`, `\"zuko_nsf\"` |\n", - "| Tune network depth or width | **2** | `posterior_nn(hidden_features=100)` |\n", - "| Add embedding networks for images/timeseries | **2** | `posterior_nn(embedding_net=my_cnn)` |\n", - "| Access specialized flow parameters | **3** | `build_nsf(num_bins=16, tail_bound=5.0)` |\n", + "| Standard workflows with good defaults | **1** | `NPE(prior)` |\n", + "| Try a different density estimator type | **2** | Switch `NSFConfig`, `ZukoMAFConfig`, `MDNConfig` |\n", + "| Tune network depth or width | **2** | `ZukoNSFConfig(hidden_features=100)` |\n", + "| Add embedding networks for images/timeseries | **2** | `ZukoNSFConfig(embedding_net=my_cnn)` |\n", + "| Access specialized flow parameters | **2** | `NSFConfig(num_bins=16, tail_bound=5.0)` |\n", + "| Reach an argument no config exposes | **3** | `partial(build_nsf, ...)` |\n", "| Implement custom network architecture | **3** | Subclass `DensityEstimator` |\n", "| Define custom loss functions or training | **4** | See advanced tutorial 18 |\n", "\n", @@ -544,8 +550,8 @@ "\n", "**4 Abstraction Levels for Density Estimator:**\n", "\n", - "- **Level 1**: Trainer classes with strings → `NPE(prior, density_estimator=\"nsf\")`\n", - "- **Level 2**: Factory functions → `posterior_nn(model=\"maf\", hidden_features=50)`\n", + "- **Level 1**: Trainer defaults → `NPE(prior)`\n", + "- **Level 2**: Config objects → `ZukoNSFConfig(hidden_features=50)`\n", "- **Level 3**: Direct builders → `build_nsf(num_bins=8, tail_bound=3.0)`\n", "- **Level 4**: Custom training → Full control (see tutorial 18)\n", "\n", From 5ac5336b5c8d879d371288f91eb0f02739266c04 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Mon, 10 Aug 2026 13:45:22 +0530 Subject: [PATCH 05/29] use configs in the custom neural nets how-to --- docs/how_to_guide/03_density_estimators.ipynb | 76 ++++++++++++------- 1 file changed, 49 insertions(+), 27 deletions(-) diff --git a/docs/how_to_guide/03_density_estimators.ipynb b/docs/how_to_guide/03_density_estimators.ipynb index 9d176fc3d..ecd22cf23 100644 --- a/docs/how_to_guide/03_density_estimators.ipynb +++ b/docs/how_to_guide/03_density_estimators.ipynb @@ -18,7 +18,7 @@ "[`nflows`](https://github.com/bayesiains/nflows/) or [`zuko`](https://github.com/probabilists/zuko). \n", "\n", "For all options, check the API reference\n", - "[here](https://sbi.readthedocs.io/en/latest/sbi.html).\n" + "[here](https://sbi.readthedocs.io/en/latest/sbi.html)." ] }, { @@ -27,20 +27,17 @@ "source": [ "## Changing the type of density estimator\n", "\n", - "One option is using one of the preconfigured density estimators by passing a string in\n", - "the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`), e.g.,\n", - "\"maf\" for a Masked Autoregressive Flow, of \"nsf\" for a Neural Spline Flow with default\n", - "hyperparameters.\n", + "The density estimator is chosen by passing a **config object** in the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`). There is one config class per model, so the choice of model is the choice of class: `MAFConfig` for a Masked Autoregressive Flow, `NSFConfig` for a Neural Spline Flow, and so on.\n", "\n", - "**New with sbi 0.23:** Note that `\"maf\"` or `\"nsf\"` correspond to `nflows` density\n", + "Note that `MAFConfig` or `NSFConfig` correspond to `nflows` density\n", "estimators. Those have proven to work well, but the `nflows` package is not maintained\n", "anymore. To use more recent and actively maintained density estimators, we tentatively\n", - "recommend using `zuko`, e.g., by passing `zuko_maf` or `zuko_nsf`. \n" + "recommend using `zuko`, e.g., `ZukoMAFConfig` or `ZukoNSFConfig`. " ] }, { "cell_type": "code", - "execution_count": 1, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ @@ -52,75 +49,100 @@ }, { "cell_type": "code", - "execution_count": 2, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ + "from sbi.neural_nets import ZukoMAFConfig\n", + "\n", "prior = BoxUniform(torch.zeros(2), torch.ones(2))\n", - "inference = NPE(prior=prior, density_estimator=\"zuko_maf\")" + "inference = NPE(prior=prior, density_estimator=ZukoMAFConfig())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "In the case of `NRE`, the argument is called `classifier`:\n" + "In the case of `NRE`, the argument is called `classifier`:" ] }, { "cell_type": "code", - "execution_count": 3, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "inference = NRE(prior=prior, classifier=\"resnet\")" + "from sbi.neural_nets import ResNetClassifierConfig\n", + "\n", + "inference = NRE(prior=prior, classifier=ResNetClassifierConfig())" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "## Changing hyperparameters of density estimators\n" + "## Changing hyperparameters of density estimators" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Alternatively, you can use a set of utils functions to configure a density estimator yourself, e.g., use a MAF with hyperparameters chosen for your problem at hand.\n", + "The hyperparameters of a model are the constructor arguments of its config, so you can tune them for the problem at hand.\n", "\n", - "Here, because we want to use N*P*E, we specifiy a neural network targeting the _posterior_ (using the utils function `posterior_nn`). In this example, we will create a neural spline flow (`'nsf'`) with `60` hidden units and `3` transform layers:\n" + "Here, because we want to use N*P*E, we pass the config to the `density_estimator` argument of `NPE`. In this example, we will create a neural spline flow (`ZukoNSFConfig`) with `60` hidden units and `3` transform layers:" ] }, { "cell_type": "code", - "execution_count": 4, + "execution_count": null, "metadata": {}, "outputs": [], "source": [ - "# For SNLE: likelihood_nn(). For SNRE: classifier_nn()\n", - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import ZukoNSFConfig\n", "\n", - "density_estimator_build_fun = posterior_nn(\n", - " model=\"zuko_nsf\", hidden_features=60, num_transforms=3\n", - ")\n", - "inference = NPE(prior=prior, density_estimator=density_estimator_build_fun)" + "density_estimator = ZukoNSFConfig(hidden_features=60, num_transforms=3)\n", + "inference = NPE(prior=prior, density_estimator=density_estimator)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "It is also possible to pass an `embedding_net` to `posterior_nn()` to automatically\n", + "A config only accepts the settings its own model has, so a setting that belongs to a different model, or a misspelled one, raises immediately instead of being ignored:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.neural_nets import ZukoMAFConfig\n", + "\n", + "try:\n", + " # `num_bins` is a spline setting, and a MAF has no splines.\n", + " ZukoMAFConfig(num_bins=8)\n", + "except TypeError as e:\n", + " print(e)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "It is also possible to pass an `embedding_net` to a config to automatically\n", "learn summary statistics from high-dimensional simulation outputs. You can find a more\n", - "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n" + "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", + "\n", + "The full list of configs, together with embedding nets, z-scoring, and how to migrate from the older string and factory-function interfaces, is in [how to configure the neural network](27_estimator_configs.ipynb)." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "## Building new density estimators from scratch\n" + "## Building new density estimators from scratch" ] }, { @@ -129,7 +151,7 @@ "source": [ "Finally, it is also possible to implement your own density estimator from scratch, e.g., including embedding nets to preprocess data, or to a density estimator architecture of your choice.\n", "\n", - "For this, the `density_estimator` argument needs to be a function that takes `theta` and `x` batches as arguments to then construct the density estimator after the first set of simulations was generated. Our factory functions in `sbi/neural_nets/factory.py` return such a function.\n", + "For this, the `density_estimator` argument needs to be a function that takes `theta` and `x` batches as arguments to then construct the density estimator after the first set of simulations was generated. This is what a config does through its `build` method, and what the factory functions in `sbi/neural_nets/factory.py` return.\n", "\n", "The returned `density_estimator` object needs to be a subclass of `DensityEstimator`, which requires to implement three methods:\n", " \n", From 567b7b510936628aab8e07ae30474723f88e02b6 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Mon, 10 Aug 2026 17:33:10 +0530 Subject: [PATCH 06/29] use configs in the density estimators tutorial --- .../03_density_estimators.ipynb | 41 +++++++++++-------- 1 file changed, 23 insertions(+), 18 deletions(-) diff --git a/docs/advanced_tutorials/03_density_estimators.ipynb b/docs/advanced_tutorials/03_density_estimators.ipynb index 969c28600..78b5a04ff 100644 --- a/docs/advanced_tutorials/03_density_estimators.ipynb +++ b/docs/advanced_tutorials/03_density_estimators.ipynb @@ -27,15 +27,16 @@ "source": [ "## Changing the type of density estimator\n", "\n", - "One option is using one of the preconfigured density estimators by passing a string in\n", - "the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`), e.g.,\n", - "\"maf\" for a Masked Autoregressive Flow, of \"nsf\" for a Neural Spline Flow with default\n", - "hyperparameters.\n", + "The density estimator is chosen by passing a **config object** in\n", + "the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`). There is one config class per model, so the choice of model is the choice of class, e.g.\n", + "`MAFConfig` for a Masked Autoregressive Flow, or `NSFConfig` for a Neural Spline Flow with default\n", + "hyperparameters. The config carries only the settings that model accepts, so an unsupported or misspelled one raises at construction instead of being ignored.\n", "\n", - "**New with sbi 0.23:** Note that `\"maf\"` or `\"nsf\"` correspond to `nflows` density\n", + "Note that `MAFConfig` or `NSFConfig` correspond to `nflows` density\n", "estimators. Those have proven to work well, but the `nflows` package is not maintained\n", "anymore. To use more recent and actively maintained density estimators, we tentatively\n", - "recommend using `zuko`, e.g., by passing `zuko_maf` or `zuko_nsf`. \n" + "recommend using `zuko`, e.g., `ZukoMAFConfig` or `ZukoNSFConfig`. \n", + "" ] }, { @@ -56,8 +57,10 @@ "metadata": {}, "outputs": [], "source": [ + "from sbi.neural_nets import ZukoMAFConfig\n", + "\n", "prior = BoxUniform(torch.zeros(2), torch.ones(2))\n", - "inference = NPE(prior=prior, density_estimator=\"zuko_maf\")" + "inference = NPE(prior=prior, density_estimator=ZukoMAFConfig())" ] }, { @@ -73,7 +76,9 @@ "metadata": {}, "outputs": [], "source": [ - "inference = NRE(prior=prior, classifier=\"resnet\")" + "from sbi.neural_nets import ResNetClassifierConfig\n", + "\n", + "inference = NRE(prior=prior, classifier=ResNetClassifierConfig())" ] }, { @@ -87,9 +92,10 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Alternatively, you can use a set of utils functions to configure a density estimator yourself, e.g., use a MAF with hyperparameters chosen for your problem at hand.\n", + "The hyperparameters of a model are the constructor arguments of its config, so you can tune them for the problem at hand.\n", "\n", - "Here, because we want to use N*P*E, we specifiy a neural network targeting the _posterior_ (using the utils function `posterior_nn`). In this example, we will create a neural spline flow (`'nsf'`) with `60` hidden units and `3` transform layers:\n" + "Here, because we want to use N*P*E, we pass the config to the `density_estimator` argument of `NPE`. In this example, we will create a neural spline flow (`ZukoNSFConfig`) with `60` hidden units and `3` transform layers:\n", + "" ] }, { @@ -98,22 +104,21 @@ "metadata": {}, "outputs": [], "source": [ - "# For SNLE: likelihood_nn(). For SNRE: classifier_nn()\n", - "from sbi.neural_nets import posterior_nn\n", + "# The same configs are used for NLE. NRE takes classifier configs instead.\n", + "from sbi.neural_nets import ZukoNSFConfig\n", "\n", - "density_estimator_build_fun = posterior_nn(\n", - " model=\"zuko_nsf\", hidden_features=60, num_transforms=3\n", - ")\n", - "inference = NPE(prior=prior, density_estimator=density_estimator_build_fun)" + "density_estimator = ZukoNSFConfig(hidden_features=60, num_transforms=3)\n", + "inference = NPE(prior=prior, density_estimator=density_estimator)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "It is also possible to pass an `embedding_net` to `posterior_nn()` to automatically\n", + "It is also possible to pass an `embedding_net` to a config to automatically\n", "learn summary statistics from high-dimensional simulation outputs. You can find a more\n", - "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n" + "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", + "" ] }, { From 8a6ddf82c82adec4a4c2634bbc7a3beb4991f336 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Tue, 11 Aug 2026 12:11:59 +0530 Subject: [PATCH 07/29] point the neural net recommendation at NSFConfig --- docs/how_to_guide/03_choose_neural_net.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/how_to_guide/03_choose_neural_net.ipynb b/docs/how_to_guide/03_choose_neural_net.ipynb index b46221a54..123ba0625 100644 --- a/docs/how_to_guide/03_choose_neural_net.ipynb +++ b/docs/how_to_guide/03_choose_neural_net.ipynb @@ -18,7 +18,7 @@ "## Explicit recommendations\n", " \n", "- When your simulation outputs are high-dimensional (e.g., images, time-series,...), we strongly recommend using an embedding net, as described [here](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", - "- For `NPE` and `NLE`, the `nsf` density estimator often outperforms the default (`maf`), but it is slower to train and may be more prone to overfitting. This can be done by setting `trainer = NPE(density_estimator=\"nsf\")`.\n", + "- For `NPE` and `NLE`, the neural spline flow often outperforms the default masked autoregressive flow, but it is slower to train and may be more prone to overfitting. This can be done by setting `trainer = NPE(density_estimator=NSFConfig())`, after `from sbi.neural_nets import NSFConfig`. See [how to configure the neural network](27_estimator_configs.ipynb) for the other models.\n", "- For best performance, optimize the hyperparameters of `sbi`." ] } From c1891d2ec85ca2f48cc05c7a26a5936c1213a9fe Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Tue, 11 Aug 2026 13:05:44 +0530 Subject: [PATCH 08/29] pass the embedding net through a config --- docs/advanced_tutorials/04_embedding_networks.ipynb | 8 ++++---- docs/how_to_guide/04_embedding_networks.ipynb | 6 +++--- 2 files changed, 7 insertions(+), 7 deletions(-) diff --git a/docs/advanced_tutorials/04_embedding_networks.ipynb b/docs/advanced_tutorials/04_embedding_networks.ipynb index 770b716b1..16d954817 100644 --- a/docs/advanced_tutorials/04_embedding_networks.ipynb +++ b/docs/advanced_tutorials/04_embedding_networks.ipynb @@ -25,7 +25,7 @@ "source": [ "```Python\n", "# import required modules\n", - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import MAFConfig\n", "\n", "# import the different choices of pre-configured embedding networks\n", "from sbi.neural_nets.embedding_nets import (\n", @@ -38,7 +38,7 @@ "embedding_net = CNNEmbedding(input_shape=(32, 32))\n", "\n", "# instantiate the conditional neural density estimator\n", - "neural_posterior = posterior_nn(model=\"maf\", embedding_net=embedding_net)\n", + "neural_posterior = MAFConfig(embedding_net=embedding_net)\n", "\n", "# setup the inference procedure with NPE\n", "inferer = NPE(prior=prior, density_estimator=neural_posterior)\n", @@ -280,10 +280,10 @@ "metadata": {}, "outputs": [], "source": [ - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import MAFConfig\n", "\n", "# instantiate the neural density estimator\n", - "neural_posterior = posterior_nn(model=\"maf\", embedding_net=embedding_net)\n", + "neural_posterior = MAFConfig(embedding_net=embedding_net)\n", "\n", "# setup the inference procedure with NPE\n", "inferer = NPE(prior=prior, density_estimator=neural_posterior)" diff --git a/docs/how_to_guide/04_embedding_networks.ipynb b/docs/how_to_guide/04_embedding_networks.ipynb index 14668f197..6067c46f5 100644 --- a/docs/how_to_guide/04_embedding_networks.ipynb +++ b/docs/how_to_guide/04_embedding_networks.ipynb @@ -36,7 +36,7 @@ "source": [ "```Python\n", "# import required modules\n", - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import MAFConfig\n", "\n", "# import the different choices of pre-configured embedding networks\n", "from sbi.neural_nets.embedding_nets import (\n", @@ -49,7 +49,7 @@ "embedding_net = CNNEmbedding(input_shape=(32, 32))\n", "\n", "# Instantiate the conditional neural density estimator\n", - "neural_posterior = posterior_nn(model=\"maf\", embedding_net=embedding_net)\n", + "neural_posterior = MAFConfig(embedding_net=embedding_net)\n", "\n", "# Setup the inference procedure with NPE\n", "trainer = NPE(density_estimator=neural_posterior)\n", @@ -93,7 +93,7 @@ "embedding_net_custom = CustomCNN()\n", "\n", "# Instantiate the conditional neural density estimator\n", - "neural_posterior = posterior_nn(model=\"maf\", embedding_net=embedding_net_custom)\n", + "neural_posterior = MAFConfig(embedding_net=embedding_net_custom)\n", "trainer = NPE(density_estimator=neural_posterior)\n", "# Continue as always...\n", "```" From b398280d3e27b07326ffd497883af78c10c4e9de Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Tue, 11 Aug 2026 15:44:21 +0530 Subject: [PATCH 09/29] use a config in the permutation invariant embedding guide --- docs/how_to_guide/08_permutation_invariant_embeddings.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/how_to_guide/08_permutation_invariant_embeddings.ipynb b/docs/how_to_guide/08_permutation_invariant_embeddings.ipynb index fc3eb2ace..f7754f0d8 100644 --- a/docs/how_to_guide/08_permutation_invariant_embeddings.ipynb +++ b/docs/how_to_guide/08_permutation_invariant_embeddings.ipynb @@ -28,7 +28,7 @@ "metadata": {}, "source": [ "```python\n", - "from sbi.neural_nets import posterior_nn\n", + "from sbi.neural_nets import MAFConfig\n", "from sbi.neural_nets.embedding_nets import FCEmbedding, PermutationInvariantEmbedding\n", "\n", "latent_dim = 10\n", @@ -45,7 +45,7 @@ " num_hiddens=10,\n", " output_dim=10,\n", ")\n", - "density_estimator = posterior_nn(\"maf\", embedding_net=embedding_net)\n", + "density_estimator = MAFConfig(embedding_net=embedding_net)\n", "trainer = NPE(density_estimator=density_estimator)\n", "```" ] From 2b073dac7a526a3f16c30758a6a10a9595c1058f Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Tue, 11 Aug 2026 16:55:00 +0530 Subject: [PATCH 10/29] drop the deprecated strings from the training how-tos --- docs/how_to_guide/07_gpu_training.ipynb | 11 +++++++---- docs/how_to_guide/23_using_pyro_with_sbi.ipynb | 3 ++- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/how_to_guide/07_gpu_training.ipynb b/docs/how_to_guide/07_gpu_training.ipynb index 5d86dff33..7d323c971 100644 --- a/docs/how_to_guide/07_gpu_training.ipynb +++ b/docs/how_to_guide/07_gpu_training.ipynb @@ -32,7 +32,9 @@ "metadata": {}, "source": [ "```python\n", - "inference = NPE(prior, device=\"cuda\", density_estimator=\"maf\")\n", + "from sbi.neural_nets import MAFConfig\n", + "\n", + "inference = NPE(prior, device=\"cuda\", density_estimator=MAFConfig())\n", "density_estimator = inference.append_simulations(theta, x, data_device=\"cpu\").train()\n", "```" ] @@ -109,15 +111,16 @@ "Whether or not you reduce your training time when training on a GPU depends on\n", "the problem at hand. We provide a couple of default density estimators for\n", "`NPE`, `NLE` and `NRE`, e.g., a mixture density network\n", - "(`density_estimator=\"mdn\"`) or a Masked Autoregressive Flow\n", - "(`density_estimator=\"maf\"`). For these default density estimators, we do **not**\n", + "(`density_estimator=MDNConfig()`) or a Masked Autoregressive Flow\n", + "(`density_estimator=MAFConfig()`). For these default density estimators, we do **not**\n", "expect a speed-up. This is because the underlying neural networks are relatively\n", "shallow and not tall, e.g., they do not have many parameters or matrix\n", "operations that benefit from being executed on the GPU.\n", "\n", "A speed-up through training on the GPU will most likely become visible when\n", "using convolutional modules in your neural networks. E.g., when passing an\n", - "embedding net for image processing like in [this example](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n" + "embedding net for image processing like in [this example](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", + "" ] } ], diff --git a/docs/how_to_guide/23_using_pyro_with_sbi.ipynb b/docs/how_to_guide/23_using_pyro_with_sbi.ipynb index 3eb5b83e5..64acc3e19 100644 --- a/docs/how_to_guide/23_using_pyro_with_sbi.ipynb +++ b/docs/how_to_guide/23_using_pyro_with_sbi.ipynb @@ -37,6 +37,7 @@ "\n", "# SBI imports\n", "from sbi.inference import NLE\n", + "from sbi.neural_nets import NSFConfig\n", "from sbi.utils.pyroutils import to_pyro_distribution\n", "\n", "pyro.clear_param_store()" @@ -122,7 +123,7 @@ "# Train Neural Likelihood Estimator from simulated data pairs\n", "trainer = NLE(\n", " prior=theta_prior,\n", - " density_estimator=\"nsf\",\n", + " density_estimator=NSFConfig(),\n", " show_progress_bars=True,\n", ")\n", "\n", From cac05eb1e27543980d90d76b0f00fd558419c4ad Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Wed, 12 Aug 2026 11:22:33 +0530 Subject: [PATCH 11/29] drop the deprecated strings from the diagnostics tutorials --- .../11_diagnostics_simulation_based_calibration.ipynb | 3 ++- ...12_iid_data_and_permutation_invariant_embeddings.ipynb | 8 +++++--- .../21_diagnostics_misspecification_checks.ipynb | 6 +++--- 3 files changed, 10 insertions(+), 7 deletions(-) diff --git a/docs/advanced_tutorials/11_diagnostics_simulation_based_calibration.ipynb b/docs/advanced_tutorials/11_diagnostics_simulation_based_calibration.ipynb index 7ca95f269..6a7f8c08f 100644 --- a/docs/advanced_tutorials/11_diagnostics_simulation_based_calibration.ipynb +++ b/docs/advanced_tutorials/11_diagnostics_simulation_based_calibration.ipynb @@ -79,6 +79,7 @@ "from sbi.analysis.plot import sbc_rank_plot\n", "from sbi.diagnostics import check_sbc, check_tarp, run_sbc, run_tarp\n", "from sbi.inference import NPE\n", + "from sbi.neural_nets import NSFConfig\n", "\n", "# Set random seed\n", "_ = torch.manual_seed(42)" @@ -175,7 +176,7 @@ ], "source": [ "# we use a mdn model to have a fast turnaround with training the NPE\n", - "inferer = NPE(prior, density_estimator=\"nsf\")\n", + "inferer = NPE(prior, density_estimator=NSFConfig())\n", "# append simulations and run training.\n", "inferer.append_simulations(theta, x).train(training_batch_size=200);" ] diff --git a/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb b/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb index 1b82d24a1..42268767d 100644 --- a/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb +++ b/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb @@ -62,6 +62,7 @@ "\n", "from sbi.analysis import pairplot\n", "from sbi.inference import NLE, NPE, simulate_for_sbi\n", + "from sbi.neural_nets import MDNConfig\n", "from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters\n", "from sbi.simulators.linear_gaussian import (\n", " linear_gaussian,\n", @@ -206,7 +207,7 @@ ], "source": [ "# Train NLE.\n", - "inferer = NLE(prior, show_progress_bars=True, density_estimator=\"mdn\")\n", + "inferer = NLE(prior, show_progress_bars=True, density_estimator=MDNConfig())\n", "theta, x = simulate_for_sbi(simulator, prior, 10000, simulation_batch_size=1000)\n", "inferer.append_simulations(theta, x).train(training_batch_size=1000);" ] @@ -380,7 +381,6 @@ "metadata": {}, "outputs": [], "source": [ - "from sbi.neural_nets import posterior_nn\n", "from sbi.neural_nets.embedding_nets import FCEmbedding, PermutationInvariantEmbedding\n", "\n", "# embedding\n", @@ -402,7 +402,9 @@ "\n", "# we choose a simple MDN as the density estimator.\n", "# NOTE: we turn off z-scoring of the data, as we used NaNs for the missing trials.\n", - "density_estimator = posterior_nn(\"mdn\", embedding_net=embedding_net, z_score_x=\"none\")" + "density_estimator = MDNConfig(\n", + " embedding_net=embedding_net, z_score_condition=\"none\"\n", + ")" ] }, { diff --git a/docs/advanced_tutorials/21_diagnostics_misspecification_checks.ipynb b/docs/advanced_tutorials/21_diagnostics_misspecification_checks.ipynb index 103fcd3c7..ca6728f9b 100644 --- a/docs/advanced_tutorials/21_diagnostics_misspecification_checks.ipynb +++ b/docs/advanced_tutorials/21_diagnostics_misspecification_checks.ipynb @@ -46,7 +46,7 @@ ")\n", "from sbi.inference import NPE\n", "from sbi.inference.trainers.marginal import MarginalTrainer\n", - "from sbi.neural_nets import MarginalNSFConfig, posterior_nn\n", + "from sbi.neural_nets import MAFConfig, MarginalNSFConfig\n", "from sbi.neural_nets.embedding_nets import FCEmbedding\n", "from sbi.utils.metrics import c2st\n", "\n", @@ -189,7 +189,7 @@ ], "source": [ "def train_npe_with_embedding(theta, x, prior, embeddding_net, **kwargs):\n", - " neural_posterior = posterior_nn(model=\"maf\", embedding_net=embeddding_net)\n", + " neural_posterior = MAFConfig(embedding_net=embeddding_net)\n", " inference = NPE(prior=prior, density_estimator=neural_posterior, **kwargs)\n", " inference = inference.append_simulations(theta, x)\n", " _ = inference.train()\n", @@ -1180,7 +1180,7 @@ "emb_net = FCEmbedding(\n", " input_dim=x_train.shape[1], output_dim=20, num_layers=4, num_hiddens=50\n", ") # minimal embedding network\n", - "neural_posterior = posterior_nn(model=\"maf\", embedding_net=emb_net)\n", + "neural_posterior = MAFConfig(embedding_net=emb_net)\n", "inference = NPE(prior=prior, density_estimator=neural_posterior)\n", "inference = inference.append_simulations(theta_train, x_train)\n", "density_estimator = inference.train()\n", From 58043e2ac73fc911b542b293ff2cb9f853bdd186 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Wed, 12 Aug 2026 14:44:11 +0530 Subject: [PATCH 12/29] drop the deprecated string from the Bayesian workflow tutorial --- docs/tutorials/01_Bayesian_workflow.ipynb | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/tutorials/01_Bayesian_workflow.ipynb b/docs/tutorials/01_Bayesian_workflow.ipynb index 649922ae6..874e5b029 100644 --- a/docs/tutorials/01_Bayesian_workflow.ipynb +++ b/docs/tutorials/01_Bayesian_workflow.ipynb @@ -373,8 +373,9 @@ "outputs": [], "source": [ "from sbi.inference import NPE\n", + "from sbi.neural_nets import NSFConfig\n", "\n", - "inference = NPE(density_estimator=\"nsf\")" + "inference = NPE(density_estimator=NSFConfig())" ] }, { From 35df04bdb1adbd6d054df6ae06755ea543ea5cf3 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Wed, 12 Aug 2026 15:30:00 +0530 Subject: [PATCH 13/29] fix the unconstrained FAQ snippet and list it in the FAQ --- docs/faq.rst | 1 + docs/faq/question_08_unconstrained.md | 13 ++++++++----- 2 files changed, 9 insertions(+), 5 deletions(-) diff --git a/docs/faq.rst b/docs/faq.rst index 59c031dfc..7e973945a 100644 --- a/docs/faq.rst +++ b/docs/faq.rst @@ -9,3 +9,4 @@ FAQ faq/question_01_leakage faq/question_02_nans faq/question_03_pickling_error + faq/question_08_unconstrained diff --git a/docs/faq/question_08_unconstrained.md b/docs/faq/question_08_unconstrained.md index aea197ee3..f2b9feed8 100644 --- a/docs/faq/question_08_unconstrained.md +++ b/docs/faq/question_08_unconstrained.md @@ -14,14 +14,17 @@ from simulations (more complex). To enable this for NPE: ```python -density_estimator_build_fun = posterior_nn( - model="zuko_nsf", +from sbi.neural_nets import ZukoNSFConfig + +density_estimator = ZukoNSFConfig( hidden_features=60, num_transforms=3, - z_score_theta="transform_to_unconstrained" # Transforms parameters to unconstrained space - x_dist=prior # For NPE, this specifies bounds for parameters (internally called 'x') + # Transforms parameters to unconstrained space. + z_score_input="transform_to_unconstrained", + # For NPE, this specifies the bounds of the parameters. + x_dist=prior, ) -inference = NPE(prior, density_estimator=density_estimator_build_fun) +inference = NPE(prior, density_estimator=density_estimator) ``` This ensures that your density estimator operates in a From 6697d46563a586b00fc0d8cc4046c093ab3bf197 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Wed, 9 Sep 2026 11:23:43 +0530 Subject: [PATCH 14/29] update the neural net config reference --- docs/api_reference/neural_nets.rst | 32 +++++++++++++++++++++++++----- 1 file changed, 27 insertions(+), 5 deletions(-) diff --git a/docs/api_reference/neural_nets.rst b/docs/api_reference/neural_nets.rst index 11089af5d..3df5ba56c 100644 --- a/docs/api_reference/neural_nets.rst +++ b/docs/api_reference/neural_nets.rst @@ -25,7 +25,6 @@ For ``NPE`` and ``NLE``, passed as ``density_estimator``. sbi.neural_nets.MAFRQSConfig sbi.neural_nets.MDNConfig sbi.neural_nets.NSFConfig - sbi.neural_nets.TabPFNConfig sbi.neural_nets.ZukoBPFConfig sbi.neural_nets.ZukoGFConfig sbi.neural_nets.ZukoMAFConfig @@ -37,6 +36,19 @@ For ``NPE`` and ``NLE``, passed as ``density_estimator``. sbi.neural_nets.ZukoUNAFConfig +Pretrained estimator configs +---------------------------- + +For ``NPE_PFN``, passed as ``density_estimator``. ``TabPFNConfig`` is not a +trainable density estimator for ``NPE`` or ``NLE``. + +.. autosummary:: + :toctree: _autosummary + :nosignatures: + + sbi.neural_nets.TabPFNConfig + + Mixed density estimator configs ------------------------------- @@ -85,16 +97,26 @@ For ``MarginalTrainer``, which fits an unconditional density. sbi.neural_nets.MarginalUNAFConfig -Vector field estimator builders -------------------------------- +Vector field estimator configs +------------------------------ -For ``FMPE`` and ``NPSE``, passed as ``vf_estimator``. +For ``FMPE`` and ``NPSE``, passed as ``vf_estimator``. The estimator config +selects flow matching or a score estimator; its ``net`` config selects the +network architecture. .. autosummary:: :toctree: _autosummary :nosignatures: - sbi.neural_nets.VectorFieldEstimatorBuilder + sbi.neural_nets.VectorFieldConfigBase + sbi.neural_nets.FlowMatchingConfig + sbi.neural_nets.ScoreConfigBase + sbi.neural_nets.VEScoreConfig + sbi.neural_nets.VPScoreConfig + sbi.neural_nets.SubVPScoreConfig + sbi.neural_nets.MLPConfig + sbi.neural_nets.AdaMLPConfig + sbi.neural_nets.TransformerConfig Factory functions From 458542f38125090936f363efa58ca8b06de4a4ab Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Wed, 9 Sep 2026 11:30:44 +0530 Subject: [PATCH 15/29] complete the estimator config and migration guide --- docs/how_to_guide/27_estimator_configs.ipynb | 110 ++++++++++++++++--- 1 file changed, 96 insertions(+), 14 deletions(-) diff --git a/docs/how_to_guide/27_estimator_configs.ipynb b/docs/how_to_guide/27_estimator_configs.ipynb index b87be414d..aff37fb71 100644 --- a/docs/how_to_guide/27_estimator_configs.ipynb +++ b/docs/how_to_guide/27_estimator_configs.ipynb @@ -11,7 +11,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Every inference method in `sbi` trains a neural network. Which network it trains, and with which settings, is described by a **config object** that you pass to the trainer.\n", + "A **config object** describes the neural estimator and its settings. Pass it to a trainer to choose and configure the network.\n", "\n", "There is one config class per model, and it carries only the settings that model actually accepts. A setting the model does not have is therefore not a field on its config, so passing it raises an error right away instead of being silently ignored.\n", "\n", @@ -19,8 +19,8 @@ "\n", "- picking a model and changing its hyperparameters\n", "- adding an embedding network and controlling z-scoring\n", - "- the configs for classifiers, mixed data, and marginal densities\n", - "- moving off the deprecated string and factory-function interfaces" + "- the configs for classifiers, mixed data, marginal densities, and vector fields\n", + "- migrating from strings and factory functions" ] }, { @@ -98,10 +98,9 @@ "| `NSFConfig` | neural spline flow (`nflows`) |\n", "| `MADEConfig` | masked autoencoder for density estimation |\n", "| `MDNConfig` | mixture density network |\n", - "| `TabPFNConfig` | TabPFN-based estimator |\n", "| `ZukoMAFConfig`, `ZukoNSFConfig`, `ZukoNCSFConfig`, `ZukoNAFConfig`, `ZukoUNAFConfig`, `ZukoBPFConfig`, `ZukoSOSPFConfig`, `ZukoNICEConfig`, `ZukoGFConfig` | the corresponding [`zuko`](https://github.com/probabilists/zuko) flows |\n", "\n", - "The `nflows` package is no longer maintained, so for new projects we tentatively recommend the `zuko` variants." + "For model recommendations, see [how to choose neural nets](03_choose_neural_net.ipynb). `NPE_A` accepts `MDNConfig` only. The pretrained `TabPFNConfig` is used with `NPE_PFN`, not with `NPE` or `NLE`." ] }, { @@ -195,7 +194,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Being frozen means a config cannot be edited after construction. To vary one setting, build a new config:" + "Config fields cannot be reassigned after construction, though contained objects such as embedding networks remain mutable. To vary one setting, build a new config:" ] }, { @@ -235,7 +234,7 @@ "source": [ "## Z-scoring\n", "\n", - "Both variables are z-scored independently by default. The standardization is built into the estimator, so you keep passing raw data at inference time.\n", + "For the trainable conditional density configs, both variables are z-scored independently by default. The standardization is built into the estimator, so you keep passing raw data at inference time.\n", "\n", "`z_score_input` applies to the modeled variable and `z_score_condition` to the variable conditioned on. Which is which depends on the method: for `NPE` the input is $\\theta$ and the condition is $x$, for `NLE` it is the other way round. Each takes `\"independent\"` (default), `\"structured\"`, or `\"none\"`:" ] @@ -254,7 +253,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "Use `\"structured\"` when the entries of a variable are not exchangeable, for example a time series, where a single mean and standard deviation across all entries would be the wrong summary. Use `\"none\"` if the data is already standardized." + "`\"independent\"` computes separate statistics for each coordinate across the batch. `\"structured\"` uses a shared mean and scale across coordinates, for example to preserve relative values within a time series or image. Use `\"none\"` to disable standardization.\n", + "\n", + "Some density configs also support `z_score_input=\"transform_to_unconstrained\"` for bounded inputs; see the [unconstrained-space FAQ](../faq/question_04_unconstrained.md)." ] }, { @@ -263,7 +264,7 @@ "source": [ "## Settings without a field\n", "\n", - "Some underlying models accept keyword arguments that have no field on the config. `extra_kwargs` forwards them, and marks at the call site that you are stepping outside the checked surface:" + "Some underlying models accept keyword arguments that have no field on the config. Use `extra_kwargs` for these library-specific options. The underlying builder must support the option; this is not a way to make an unsupported setting work:" ] }, { @@ -351,20 +352,23 @@ "metadata": {}, "outputs": [], "source": [ + "from sbi.inference import MNLE\n", "from sbi.neural_nets import MixedConfig\n", "\n", "config = MixedConfig(\n", " continuous=ZukoNSFConfig(hidden_features=64, num_transforms=4),\n", " discrete_hidden_features=32,\n", ")\n", - "print(config)" + "trainer = MNLE(prior=prior, density_estimator=config)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ - "Note that z-scoring of the modeled variable is a property of the continuous model, so it lives on the nested config, while `z_score_condition` stays on `MixedConfig`." + "The modeled variable is $x$ for MNLE and $\\theta$ for MNPE. Put its continuous columns first and discrete columns last. `MixedConfig()` defaults to an nflows NSF continuous component with `tail_bound=10.0`.\n", + "\n", + "Set input z-scoring on `continuous.z_score_input`. Set condition z-scoring and the condition embedding on `MixedConfig`, not on its nested continuous config. Put library-specific continuous options in `continuous.extra_kwargs`; top-level `MixedConfig.extra_kwargs` is unsupported." ] }, { @@ -390,6 +394,49 @@ ")" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Vector fields for FMPE and NPSE\n", + "\n", + "Pass an estimator config as `vf_estimator`. `FlowMatchingConfig` is for FMPE; `VEScoreConfig`, `VPScoreConfig`, and `SubVPScoreConfig` select NPSE's SDE. Choose the architecture separately through `net`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "from sbi.inference import FMPE, NPSE\n", + "from sbi.neural_nets import (\n", + " FlowMatchingConfig,\n", + " MLPConfig,\n", + " TransformerConfig,\n", + " VPScoreConfig,\n", + ")\n", + "\n", + "flow_config = FlowMatchingConfig(net=MLPConfig(hidden_features=64, num_layers=3))\n", + "flow_trainer = FMPE(prior=prior, vf_estimator=flow_config)\n", + "\n", + "score_config = VPScoreConfig(\n", + " net=TransformerConfig(hidden_features=64, num_heads=4), beta_max=20.0\n", + ")\n", + "score_trainer = NPSE(prior=prior, vf_estimator=score_config)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "The network configs are `MLPConfig`, `AdaMLPConfig`, and `TransformerConfig`. `TransformerConfig(is_x_emb_seq=True)` selects cross-attention for sequence conditions. Both `FlowMatchingConfig()` and the default `VEScoreConfig()` use `MLPConfig()`.\n", + "\n", + "Keep `embedding_net`, `z_score_input`, and `z_score_condition` on the estimator config. Estimator options without fields go in its `extra_kwargs`; network options go in `net.extra_kwargs`. Do not pass `sde_type` to NPSE together with a config: the config class already selects it.\n", + "\n", + "For custom networks, pass a `VectorFieldNet` instance as `net`. See the [vector-field tutorial](../advanced_tutorials/19_vector_field_methods.ipynb) for examples and noise schedules." + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -413,7 +460,24 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The first argument is the modeled variable and the second the one conditioned on, matching `z_score_input` and `z_score_condition`." + "The first argument is the modeled variable and the second the one conditioned on, matching `z_score_input` and `z_score_condition`. Marginal configs take only one batch: `MarginalNSFConfig().build(x)`.\n", + "\n", + "For a custom estimator, pass a build function to the trainer. Unlike `config.build`, the trainer calls this function with `(batch_theta, batch_x)` for both NPE and NLE. Return an estimator implementing the required interface, not an arbitrary `torch.nn.Module`:" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "def build_likelihood(batch_theta, batch_x):\n", + " return MAFConfig(hidden_features=32).build(\n", + " batch_input=batch_x, batch_condition=batch_theta\n", + " )\n", + "\n", + "\n", + "trainer = NLE(prior=prior, density_estimator=build_likelihood)" ] }, { @@ -433,8 +497,26 @@ "| `posterior_nn(model=\"nsf\", hidden_features=64)` | `NSFConfig(hidden_features=64)` |\n", "| `likelihood_nn(model=\"maf\", num_transforms=8)` | `MAFConfig(num_transforms=8)` |\n", "| `classifier_nn(model=\"mlp\", hidden_features=64)` | `MLPClassifierConfig(hidden_features=64)` |\n", + "| `MNLE(prior, density_estimator=\"mnle\")` | `MNLE(prior, density_estimator=MixedConfig())` |\n", + "| `MNPE(prior, density_estimator=\"mnpe\")` | `MNPE(prior, density_estimator=MixedConfig())` |\n", + "| `MarginalTrainer(density_estimator=ZukoFlowType.NSF)` | `MarginalTrainer(density_estimator=MarginalNSFConfig())` |\n", + "| `posterior_flow_nn(model=\"mlp\", hidden_features=64)` | `FlowMatchingConfig(net=MLPConfig(hidden_features=64))` |\n", + "| `posterior_score_nn(model=\"mlp\", sde_type=\"vp\")` | `VPScoreConfig(net=MLPConfig())` |\n", + "\n", + "Most model-specific settings keep their names, but z-scoring names depend on the variable's role:\n", + "\n", + "| Factory argument | Config field |\n", + "|---|---|\n", + "| `posterior_nn(z_score_theta=...)` | `z_score_input` |\n", + "| `posterior_nn(z_score_x=...)` | `z_score_condition` |\n", + "| `likelihood_nn(z_score_x=...)` | `z_score_input` |\n", + "| `likelihood_nn(z_score_theta=...)` | `z_score_condition` |\n", + "| `classifier_nn(z_score_theta=..., z_score_x=...)` | `z_score_input`, `z_score_condition` |\n", + "| `marginal_nn(z_score_x=...)` | `z_score_input` |\n", + "| `posterior_flow_nn` / `posterior_score_nn`: `z_score_theta`, `z_score_x` | `z_score_input`, `z_score_condition` |\n", + "| VF factories: `t_embedding_dim` | `net.time_embedding_dim` |\n", "\n", - "The model name moves from a string argument into the class name, and the remaining keyword arguments carry over unchanged." + "For marginal spline configs, use `bins` instead of the factory's `num_bins`. Use `\"none\"` to disable z-scoring in configs rather than the legacy `None`." ] }, { @@ -458,7 +540,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "One difference worth knowing when you migrate: the factory functions accept settings the chosen model ignores, whereas the config rejects them. If a call that used to run now raises a `TypeError`, that setting was not reaching the network before either." + "Configs reject misspelled or inapplicable fields at construction. The density, classifier, mixed, and vector-field factories now also reject supplied non-default settings that their selected model does not use. Genuinely unknown factory keywords still warn and are forwarded for compatibility. With configs, use `extra_kwargs` only for options supported by the underlying library." ] }, { From 8283917553a31190c86c15c3a5ef6320a166eb7f Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Thu, 10 Sep 2026 08:12:11 +0530 Subject: [PATCH 16/29] clarify custom density estimator builders --- docs/advanced_tutorials/03_density_estimators.ipynb | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/docs/advanced_tutorials/03_density_estimators.ipynb b/docs/advanced_tutorials/03_density_estimators.ipynb index a8cd39a24..433120dac 100644 --- a/docs/advanced_tutorials/03_density_estimators.ipynb +++ b/docs/advanced_tutorials/03_density_estimators.ipynb @@ -18,7 +18,7 @@ "[`nflows`](https://github.com/bayesiains/nflows/) or [`zuko`](https://github.com/probabilists/zuko). \n", "\n", "For all options, check the API reference\n", - "[here](https://sbi.readthedocs.io/en/latest/api_reference.html#neural-nets)." + "[here](../api_reference/neural_nets.rst)." ] }, { @@ -136,7 +136,7 @@ "\n", "For this, the `density_estimator` argument needs to be a function that takes `theta` and `x` batches as arguments to then construct the density estimator after the first set of simulations was generated. Our factory functions in `sbi/neural_nets/factory.py` return such a function.\n", "\n", - "The returned `density_estimator` object needs to be a subclass of [`DensityEstimator`](https://github.com/sbi-dev/sbi/blob/1928f018fa08bb0c5309a34d8e95b9f2916b20a5/sbi/neural_nets/estimators/base.py#L11), which requires to implement three methods:\n", + "The returned estimator must subclass `ConditionalDensityEstimator` from `sbi.neural_nets.estimators` and implement three methods:\n", " \n", "- `log_prob(input, condition, **kwargs)`: Return the log probabilities of the inputs given a condition or multiple i.e. batched conditions.\n", "- `loss(input, condition, **kwargs)`: Return the loss for training the density estimator.\n", From 851c6407507d73cf188c98ac679fa5f1b632c268 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Thu, 10 Sep 2026 09:00:28 +0530 Subject: [PATCH 17/29] fix the density estimator tutorial references --- docs/how_to_guide/03_density_estimators.ipynb | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/how_to_guide/03_density_estimators.ipynb b/docs/how_to_guide/03_density_estimators.ipynb index 83b175c5a..726cf047e 100644 --- a/docs/how_to_guide/03_density_estimators.ipynb +++ b/docs/how_to_guide/03_density_estimators.ipynb @@ -18,7 +18,7 @@ "[`nflows`](https://github.com/bayesiains/nflows/) or [`zuko`](https://github.com/probabilists/zuko). \n", "\n", "For all options, check the API reference\n", - "[here](https://sbi.readthedocs.io/en/latest/api_reference.html).\n" + "[here](../api_reference/neural_nets.rst)." ] }, { @@ -151,15 +151,15 @@ "source": [ "Finally, it is also possible to implement your own density estimator from scratch, e.g., including embedding nets to preprocess data, or to a density estimator architecture of your choice.\n", "\n", - "For this, the `density_estimator` argument needs to be a function that takes `theta` and `x` batches as arguments to then construct the density estimator after the first set of simulations was generated. This is what a config does through its `build` method, and what the factory functions in `sbi/neural_nets/factory.py` return.\n", + "For this, the `density_estimator` argument needs to be a function that takes `theta` and `x` batches and returns an estimator. The trainer calls it after receiving simulations. Unlike this callable interface, `config.build` takes the modeled variable first: `build(theta, x)` for NPE and `build(x, theta)` for NLE.\n", "\n", - "The returned `density_estimator` object needs to be a subclass of `DensityEstimator`, which requires to implement three methods:\n", + "The returned estimator must subclass `ConditionalDensityEstimator` from `sbi.neural_nets.estimators` and implement three methods:\n", " \n", "- `log_prob(input, condition, **kwargs)`: Return the log probabilities of the inputs given a condition or multiple i.e. batched conditions.\n", "- `loss(input, condition, **kwargs)`: Return the loss for training the density estimator.\n", "- `sample(sample_shape, condition, **kwargs)`: Return samples from the density estimator.\n", "\n", - "See more information on the [Reference API page](https://sbi.readthedocs.io/en/latest/api_reference.html)." + "See the [training interface tutorial](../advanced_tutorials/18_training_interface.ipynb) for custom training loops." ] } ], From 5ecdffc8aaa3d8b6502c5a557e25210c570e0f54 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Thu, 10 Sep 2026 09:26:35 +0530 Subject: [PATCH 18/29] fix custom builder and extra kwargs examples --- docs/how_to_guide/24_abstraction_levels.ipynb | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/docs/how_to_guide/24_abstraction_levels.ipynb b/docs/how_to_guide/24_abstraction_levels.ipynb index 350e4d344..1fcff50ba 100644 --- a/docs/how_to_guide/24_abstraction_levels.ipynb +++ b/docs/how_to_guide/24_abstraction_levels.ipynb @@ -172,7 +172,7 @@ "\n", "**Use case**: Reach a build-function argument that no config exposes\n", "\n", - "Pass any callable taking the two training batches. The `build_*` functions are such callables once their settings are bound, which gives access to their full signature." + "Pass a callable taking `(batch_theta, batch_x)` and returning an estimator. The example below uses a low-level builder directly for NPE. For NLE, wrap the builder to pass `batch_x` as the modeled input and `batch_theta` as its condition; see the [config guide](27_estimator_configs.ipynb)." ] }, { @@ -212,9 +212,9 @@ "**Key features**:\n", "- Direct access to all builder parameters\n", "- Maximum flexibility for architecture design\n", - "- Can implement fully custom architectures by subclassing `DensityEstimator`\n", + "- Can implement custom estimators by subclassing `ConditionalDensityEstimator`\n", "\n", - "For a setting the config does not have but the build function does, `extra_kwargs` on the config is usually the smaller step: `ZukoMAFConfig(extra_kwargs={\"randperm\": True})`." + "For a setting the config does not have but the build function does, use `extra_kwargs`, for example `ZukoNSFConfig(extra_kwargs={\"randperm\": True})`." ] }, { @@ -518,7 +518,7 @@ "| Add embedding networks for images/timeseries | **2** | `ZukoNSFConfig(embedding_net=my_cnn)` |\n", "| Access specialized flow parameters | **2** | `NSFConfig(num_bins=16, tail_bound=5.0)` |\n", "| Reach an argument no config exposes | **3** | `partial(build_nsf, ...)` |\n", - "| Implement custom network architecture | **3** | Subclass `DensityEstimator` |\n", + "| Implement custom network architecture | **3** | Subclass `ConditionalDensityEstimator` |\n", "| Define custom loss functions or training | **4** | See advanced tutorial 18 |\n", "\n", "**Rule of thumb**: Start with Level 1. Move to higher levels only when you need specific control." From 278b9b3a29d566fc5d64f6ef4cf5d7894e76ffeb Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 11 Sep 2026 15:41:10 +0530 Subject: [PATCH 19/29] clarify unconstrained transform support --- docs/faq/question_04_unconstrained.md | 25 +++++++++---------------- 1 file changed, 9 insertions(+), 16 deletions(-) diff --git a/docs/faq/question_04_unconstrained.md b/docs/faq/question_04_unconstrained.md index fa6d33a74..a6cd0f8fe 100644 --- a/docs/faq/question_04_unconstrained.md +++ b/docs/faq/question_04_unconstrained.md @@ -2,16 +2,15 @@ Yes. If posterior samples leak outside the prior bounds (see [posterior samples outside the prior support](question_01_leakage.md)) and you have -ruled out simulator issues, you can train your density or ratio estimator in an +ruled out simulator issues, you can train a supported density estimator in an unbounded space using a logit transformation: - **For NPE**: The transformation maps bounded parameters θ to unbounded space before training, then applies the inverse (sigmoid) after training to ensure posterior samples stay within prior bounds. -- **For NLE/NRE**: The transformation would need to map bounded -data x to unbounded space, which requires estimating data bounds -from simulations (more complex). +- **For NLE**: The modeled variable is the data x. Supply a distribution with +the appropriate data support as `x_dist`, rather than the parameter prior. To enable this for NPE: @@ -27,9 +26,7 @@ prior = BoxUniform(low=torch.zeros(2), high=torch.ones(2)) density_estimator = ZukoNSFConfig( hidden_features=60, num_transforms=3, - # Transforms parameters to unconstrained space. z_score_input="transform_to_unconstrained", - # For NPE, this specifies the bounds of the parameters. x_dist=prior, ) inference = NPE(prior, density_estimator=density_estimator) @@ -39,16 +36,12 @@ This ensures that your density estimator operates in a transformed space where it respects prior bounds, improving the efficiency of rejection sampling. -Note: The `x_dist=prior` might seem confusing - internally, -sbi uses generic `x,y` notation where for NPE, `x` represents -parameters (θ) and `y` represents data. -This is why we pass the prior as `x_dist`. +Despite its name, `x_dist` describes the support of the modeled variable. +For NPE, that variable is θ, so pass the prior. Important: -- This transformation is currently supported by the conditional zuko density -estimators (for example `zuko_maf` and `zuko_nsf`) and by `mdn`. The nflows-based -estimators reject it, and so do the unconditional (marginal) flows. -- For **NLE/NRE**, setting up this transformation is more -complex as it requires estimating bounds for the simulated data -rather than using prior bounds. +- Zuko density configs and `MDNConfig` support this transformation. +- The nflows density configs, marginal configs, classifier configs, and +vector-field configs do not offer `"transform_to_unconstrained"` as an input +z-scoring mode. From 0d83349344a619b131af9fdb26fff9275b9491ba Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sun, 13 Sep 2026 06:55:31 +0530 Subject: [PATCH 20/29] use configs in the vector field options guide --- .../25_choosing_vector_field_options.ipynb | 43 +++++++++++-------- 1 file changed, 26 insertions(+), 17 deletions(-) diff --git a/docs/how_to_guide/25_choosing_vector_field_options.ipynb b/docs/how_to_guide/25_choosing_vector_field_options.ipynb index af3fecb96..a3cfa53f8 100644 --- a/docs/how_to_guide/25_choosing_vector_field_options.ipynb +++ b/docs/how_to_guide/25_choosing_vector_field_options.ipynb @@ -43,7 +43,7 @@ "| **Sampling quality** | Good for smooth posteriors | More robust for complex posteriors |\n", "| **IID posteriors** | Supported (via bridge) | Supported (native) |\n", "| **Guidance** | Supported (via bridge) | Supported (native) |\n", - "| **EDM schedules** | Not applicable | Supported for `sde_type=\"ve\"` |" + "| **EDM schedules** | Not applicable | Supported by `VEScoreConfig` |" ] }, { @@ -53,11 +53,13 @@ "source": [ "## Choosing an architecture\n", "\n", - "Use the default `\"mlp\"` unless you have a specific reason:\n", - "- **< 20 parameter dims**: `\"mlp\"` (default) — fast, reliable.\n", - "- **20–100 parameter dims or complex posteriors**: try `\"ada_mlp\"` — uses FiLM-style time conditioning.\n", - "- **> 100 parameter dims or structured data**: try `\"transformer\"` — benefits strongly from GPU.\n", - "- **Variable-length observation sequences** (e.g., time series): use `\"transformer_cross_attn\"` with an appropriate embedding net.\n", + "Use the default `MLPConfig()` unless you have a specific reason:\n", + "- **< 20 parameter dims**: `MLPConfig` (default) — fast, reliable.\n", + "- **20–100 parameter dims or complex posteriors**: try `AdaMLPConfig` — uses FiLM-style time conditioning.\n", + "- **> 100 parameter dims or structured data**: try `TransformerConfig` — benefits strongly from GPU.\n", + "- **Sequence conditions**: `TransformerConfig(is_x_emb_seq=True)` selects cross-attention.\n", + "\n", + "Nest the network config in `FlowMatchingConfig` or a score config, then pass it as the trainer's `vf_estimator`.\n", "\n", "See [How to choose neural nets](https://sbi.readthedocs.io/en/latest/how_to_guide/03_choose_neural_net.html) and the [tutorial Section 4](../advanced_tutorials/19_vector_field_methods.ipynb) for details." ] @@ -69,13 +71,21 @@ "metadata": {}, "outputs": [], "source": [ - "from sbi.neural_nets import posterior_flow_nn, posterior_score_nn\n", + "from sbi.neural_nets import (\n", + " AdaMLPConfig,\n", + " FlowMatchingConfig,\n", + " TransformerConfig,\n", + " VEScoreConfig,\n", + " VPScoreConfig,\n", + ")\n", "\n", "# FMPE with transformer\n", - "net = posterior_flow_nn(model=\"transformer\", num_layers=2, num_heads=2, hidden_features=64)\n", + "flow_config = FlowMatchingConfig(\n", + " net=TransformerConfig(num_layers=2, num_heads=2, hidden_features=64)\n", + ")\n", "\n", "# NPSE with ada_mlp\n", - "net = posterior_score_nn(model=\"ada_mlp\", sde_type=\"ve\", hidden_features=128, num_layers=6)" + "score_config = VEScoreConfig(net=AdaMLPConfig(hidden_features=128, num_layers=6))" ] }, { @@ -85,11 +95,11 @@ "source": [ "## Choosing an SDE type (NPSE only)\n", "\n", - "- `\"ve\"` (default): Best general choice. Supports EDM-style noise schedules.\n", - "- `\"vp\"`: Can work better for low-dimensional problems.\n", - "- `\"subvp\"`: Experimental, tighter variance bounds.\n", + "- `VEScoreConfig` (default): Best general choice. Supports EDM-style noise schedules.\n", + "- `VPScoreConfig`: Can work better for low-dimensional problems.\n", + "- `SubVPScoreConfig`: Experimental, tighter variance bounds.\n", "\n", - "If unsure, use `\"ve\"`." + "If unsure, use `VEScoreConfig()`. The config class selects the SDE, so do not also pass `sde_type` to NPSE." ] }, { @@ -99,7 +109,7 @@ "source": [ "## Tuning noise schedules\n", "\n", - "All SDE types have tunable noise range parameters: `sigma_min`/`sigma_max` for VE, `beta_min`/`beta_max` for VP/SubVP. These can be passed to `posterior_score_nn()`.\n", + "All SDE types have tunable noise range parameters: `sigma_min`/`sigma_max` on `VEScoreConfig`, and `beta_min`/`beta_max` on `VPScoreConfig` and `SubVPScoreConfig`.\n", "\n", "Additionally, VE supports EDM-style non-uniform time sampling. Consider these when:\n", "- Training loss is very noisy.\n", @@ -117,14 +127,13 @@ "outputs": [], "source": [ "# VE with EDM-style schedules\n", - "net_ve = posterior_score_nn(\n", - " model=\"mlp\", sde_type=\"ve\",\n", + "config_ve = VEScoreConfig(\n", " train_schedule=\"lognormal\", solve_schedule=\"power_law\",\n", " sigma_min=1e-3, sigma_max=15.0,\n", ")\n", "\n", "# VP: tune beta range\n", - "net_vp = posterior_score_nn(model=\"mlp\", sde_type=\"vp\", beta_min=0.01, beta_max=20.0)" + "config_vp = VPScoreConfig(beta_min=0.01, beta_max=20.0)" ] }, { From f33ad1669b2abdf1422389750c8d2f8195994727 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Sun, 13 Sep 2026 07:21:19 +0530 Subject: [PATCH 21/29] use configs in the vector field tutorial --- .../19_vector_field_methods.ipynb | 56 +++++++++---------- 1 file changed, 26 insertions(+), 30 deletions(-) diff --git a/docs/advanced_tutorials/19_vector_field_methods.ipynb b/docs/advanced_tutorials/19_vector_field_methods.ipynb index 67e94bb7d..d12d311d8 100644 --- a/docs/advanced_tutorials/19_vector_field_methods.ipynb +++ b/docs/advanced_tutorials/19_vector_field_methods.ipynb @@ -69,7 +69,12 @@ "\n", "from sbi.analysis import pairplot\n", "from sbi.inference import FMPE, NPSE\n", - "from sbi.neural_nets import posterior_flow_nn, posterior_score_nn\n", + "from sbi.neural_nets import (\n", + " FlowMatchingConfig,\n", + " MLPConfig,\n", + " TransformerConfig,\n", + " VEScoreConfig,\n", + ")\n", "from sbi.utils import BoxUniform" ] }, @@ -206,7 +211,7 @@ "\n", "Note that only the **single-round** version of NPSE is implemented currently.\n", "\n", - "In `sbi`, the `sde_type` parameter defines whether the forward diffusion process has a noising schedule that is Variance Exploding (`\"ve\"`, i.e., [SMLD](https://proceedings.neurips.cc/paper/2019/hash/3001ef257407d5a371a96dcd947c7d93-Abstract.html)), Variance Preserving (`\"vp\"`, i.e., [DDPM](https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html)), or sub-Variance Preserving (`\"subvp\"`)." + "In `sbi`, choose the forward diffusion process with `VEScoreConfig` (variance exploding, i.e., [SMLD](https://proceedings.neurips.cc/paper/2019/hash/3001ef257407d5a371a96dcd947c7d93-Abstract.html)), `VPScoreConfig` (variance preserving, i.e., [DDPM](https://proceedings.neurips.cc/paper/2020/hash/4c5bcfec8584af0d967f1ab10179ca4b-Abstract.html)), or `SubVPScoreConfig` (sub-variance preserving)." ] }, { @@ -235,7 +240,7 @@ ], "source": [ "# Minimal NPSE example\n", - "npse_trainer = NPSE(prior, sde_type=\"ve\")\n", + "npse_trainer = NPSE(prior, vf_estimator=VEScoreConfig())\n", "npse_trainer.append_simulations(theta, x).train()\n", "npse_posterior = npse_trainer.build_posterior()\n", "samples_npse = npse_posterior.sample((num_posterior_samples,), x=x_o)\n", @@ -259,12 +264,12 @@ "However, performance can be highly dependent on the suitability of the chosen neural network for the task at hand. For instance, diffusion models achieve state-of-the-art generative performance on images only when paired with a carefully designed U-Net-like architecture.\n", "\n", "For both FMPE and NPSE, we provide a selection of neural networks tailored to the task:\n", - "- `\"mlp\"` and `\"ada_mlp\"` (MLP with adaptive layer norm for conditioning) — MLP variants that inject time information at each layer, as is common in diffusion architectures.\n", - "- `\"transformer\"` and `\"transformer_cross_attn\"` — scalable diffusion transformers (e.g., as in [Peebles & Xie, 2023](https://arxiv.org/abs/2212.09748)). The cross-attention variant supports arbitrary sequence lengths for conditioning.\n", + "- `MLPConfig` and `AdaMLPConfig` (MLP with adaptive layer norm for conditioning) — MLP variants that inject time information at each layer, as is common in diffusion architectures.\n", + "- `TransformerConfig` — scalable diffusion transformers (e.g., as in [Peebles & Xie, 2023](https://arxiv.org/abs/2212.09748)). Set `is_x_emb_seq=True` to select cross-attention for sequence conditions.\n", "\n", "For certain data types, such as image-like inputs, it may still be beneficial to use a more specialized architecture.\n", "\n", - "Use `posterior_flow_nn()` for FMPE and `posterior_score_nn()` for NPSE to configure the network." + "Pass the network config as `net` inside the estimator config. Set `embedding_net`, `z_score_input`, and `z_score_condition` on the estimator config. For both methods, the input is $\\theta$ and the condition is $x$. See the [config guide](../how_to_guide/27_estimator_configs.ipynb) for migration from the legacy factories." ] }, { @@ -283,13 +288,10 @@ ], "source": [ "# FMPE with a transformer architecture\n", - "net_builder = posterior_flow_nn(\n", - " model=\"transformer\",\n", - " num_layers=2,\n", - " num_heads=2,\n", - " hidden_features=64,\n", + "config = FlowMatchingConfig(\n", + " net=TransformerConfig(num_layers=2, num_heads=2, hidden_features=64),\n", ")\n", - "trainer = FMPE(prior, vf_estimator=net_builder)\n", + "trainer = FMPE(prior, vf_estimator=config)\n", "estimator = trainer.append_simulations(theta, x).train(\n", " training_batch_size=200, learning_rate=5e-4\n", ")" @@ -310,14 +312,11 @@ } ], "source": [ - "# NPSE with custom network configuration via posterior_score_nn\n", - "net_builder = posterior_score_nn(\n", - " model=\"mlp\",\n", - " sde_type=\"ve\",\n", - " hidden_features=128,\n", - " num_layers=6,\n", + "# NPSE with a custom MLP configuration\n", + "config = VEScoreConfig(\n", + " net=MLPConfig(hidden_features=128, num_layers=6),\n", ")\n", - "trainer = NPSE(prior, vf_estimator=net_builder)\n", + "trainer = NPSE(prior, vf_estimator=config)\n", "estimator = trainer.append_simulations(theta, x).train()" ] }, @@ -326,10 +325,10 @@ "id": "cell-11", "metadata": {}, "source": [ - "### Custom networks via the `VectorFieldNet` protocol\n", + "### Custom networks via `VectorFieldNet`\n", "\n", "You can use your custom network for estimating the vector field, essentially any\n", - "`torch.nn.Module` that follows the `VectorFieldNet` protocol: it must accept `(theta, x,\n", + "`torch.nn.Module` implementing the `VectorFieldNet` interface: it must accept `(theta, x,\n", "t)` and return a tensor with the same shape as `theta`." ] }, @@ -366,9 +365,8 @@ " h = torch.cat([theta, x, t[..., None]], dim=-1)\n", " return self.layers(h)\n", "\n", - "# Wrap in the factory function (adds z-scoring).\n", - "net_builder = posterior_flow_nn(model=CustomNet())\n", - "trainer = FMPE(prior, vf_estimator=net_builder)\n", + "config = FlowMatchingConfig(net=CustomNet())\n", + "trainer = FMPE(prior, vf_estimator=config)\n", "estimator = trainer.append_simulations(theta, x).train()" ] }, @@ -390,11 +388,11 @@ "- **VE**: `sigma_min` / `sigma_max` (default: 1e-4, 10.0) — controls the range of noise added during diffusion.\n", "- **VP / SubVP**: `beta_min` / `beta_max` (default: 0.01, 10.0) — controls the linear noise schedule.\n", "\n", - "These can be passed as `**kwargs` to `posterior_score_nn()` and are validated by `ScoreEstimatorConfig`.\n", + "Set these on `VEScoreConfig`, `VPScoreConfig`, or `SubVPScoreConfig`.\n", "\n", "### EDM-style time sampling schedules (VE only)\n", "\n", - "For NPSE with `sde_type=\"ve\"`, you can use EDM-style schedules from [Karras et al. 2022](https://arxiv.org/abs/2206.00364):\n", + "For NPSE with `VEScoreConfig`, you can use EDM-style schedules from [Karras et al. 2022](https://arxiv.org/abs/2206.00364):\n", "- `train_schedule=\"lognormal\"`: Concentrates training on intermediate noise levels where the score is most informative.\n", "- `solve_schedule=\"power_law\"`: Concentrates sampling steps near low noise levels for sharper samples." ] @@ -415,9 +413,7 @@ ], "source": [ "# NPSE with EDM-style noise schedules\n", - "net_builder = posterior_score_nn(\n", - " model=\"mlp\",\n", - " sde_type=\"ve\",\n", + "config = VEScoreConfig(\n", " # EDM-style training schedule\n", " train_schedule=\"lognormal\",\n", " lognormal_mean=-1.2,\n", @@ -429,7 +425,7 @@ " sigma_min=1e-3, # default: 1e-4\n", " sigma_max=15.0, # default: 10.0; increased for broader prior\n", ")\n", - "trainer = NPSE(prior, vf_estimator=net_builder)\n", + "trainer = NPSE(prior, vf_estimator=config)\n", "trainer.append_simulations(theta, x).train()\n", "posterior = trainer.build_posterior()" ] From 4fc3dbed28cf8f4e8754be2cc280f980644d9036 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 08:57:36 +0530 Subject: [PATCH 22/29] fix import order in the iid data tutorial --- .../12_iid_data_and_permutation_invariant_embeddings.ipynb | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb b/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb index 42268767d..e25bde47c 100644 --- a/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb +++ b/docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb @@ -62,8 +62,8 @@ "\n", "from sbi.analysis import pairplot\n", "from sbi.inference import NLE, NPE, simulate_for_sbi\n", - "from sbi.neural_nets import MDNConfig\n", "from sbi.inference.posteriors.posterior_parameters import MCMCPosteriorParameters\n", + "from sbi.neural_nets import MDNConfig\n", "from sbi.simulators.linear_gaussian import (\n", " linear_gaussian,\n", " true_posterior_linear_gaussian_mvn_prior,\n", From 9fa6cee9879a21c949ce403ac641d4e1f5f4ee92 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 08:58:01 +0530 Subject: [PATCH 23/29] point the API overview to per-model configs --- docs/api_reference.rst | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/docs/api_reference.rst b/docs/api_reference.rst index d2d5e23e4..00989dbec 100644 --- a/docs/api_reference.rst +++ b/docs/api_reference.rst @@ -39,6 +39,10 @@ Prior and simulator Neural nets ----------- +Use the per-model configs listed in :doc:`api_reference/neural_nets` to select +and configure an estimator. The following factory functions remain available +for backwards compatibility. + .. autosummary:: :nosignatures: From 18e02f7689ecfcc2c5e24f0333c5a3315963c4bf Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 08:58:12 +0530 Subject: [PATCH 24/29] add the estimator config guide to the docs index --- docs/llms.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/llms.txt b/docs/llms.txt index 9a96386c3..d5d81a229 100644 --- a/docs/llms.txt +++ b/docs/llms.txt @@ -25,6 +25,7 @@ expected coverage, TARP, L-C2ST). - [Summary Statistics](https://sbi.readthedocs.io/en/latest/how_to_guide/01_crafting_summary_statistics.html): Hand-crafted features for high-dim data - [Multi-Round Inference](https://sbi.readthedocs.io/en/latest/how_to_guide/02_multiround_inference.html): Sequential methods (SNPE, SNLE, SNRE) - [Density Estimators](https://sbi.readthedocs.io/en/latest/how_to_guide/03_density_estimators.html): Choosing and configuring neural networks +- [Estimator Configs](https://sbi.readthedocs.io/en/latest/how_to_guide/27_estimator_configs.html): Per-model configs, z-scoring, embeddings, and migration from factories - [Embedding Networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html): FCEmbedding, CNNEmbedding, custom nn.Module - [Choosing Inference Method](https://sbi.readthedocs.io/en/latest/how_to_guide/06_choosing_inference_method.html): NPE vs NLE vs NRE decision guide - [Sampler Interface](https://sbi.readthedocs.io/en/latest/how_to_guide/09_sampler_interface.html): Choosing MCMC, rejection, VI or importance sampling for NLE/NRE posteriors From b7b0d4f008cf6437d28d4a1e7b6cd077828409a9 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 09:14:50 +0530 Subject: [PATCH 25/29] clarify the mixed estimator description --- docs/api_reference/neural_nets.rst | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/docs/api_reference/neural_nets.rst b/docs/api_reference/neural_nets.rst index 3df5ba56c..7e2a4b906 100644 --- a/docs/api_reference/neural_nets.rst +++ b/docs/api_reference/neural_nets.rst @@ -52,7 +52,8 @@ trainable density estimator for ``NPE`` or ``NLE``. Mixed density estimator configs ------------------------------- -For ``MNPE`` and ``MNLE``, where part of the data is discrete. +For ``MNPE`` and ``MNLE``, where the modeled variable has continuous and +discrete components. .. autosummary:: :toctree: _autosummary From d76f726dfa175c46686c2003cb24dc3934032a91 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 09:15:16 +0530 Subject: [PATCH 26/29] tighten config guide explain extra kwargs warnings --- docs/how_to_guide/27_estimator_configs.ipynb | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/docs/how_to_guide/27_estimator_configs.ipynb b/docs/how_to_guide/27_estimator_configs.ipynb index aff37fb71..e4b6c3a27 100644 --- a/docs/how_to_guide/27_estimator_configs.ipynb +++ b/docs/how_to_guide/27_estimator_configs.ipynb @@ -13,7 +13,7 @@ "source": [ "A **config object** describes the neural estimator and its settings. Pass it to a trainer to choose and configure the network.\n", "\n", - "There is one config class per model, and it carries only the settings that model actually accepts. A setting the model does not have is therefore not a field on its config, so passing it raises an error right away instead of being silently ignored.\n", + "Each model has its own config class. Unsupported or misspelled fields raise an error at construction.\n", "\n", "This guide covers:\n", "\n", @@ -107,7 +107,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "`NLE` estimates the likelihood rather than the posterior, but it takes the same configs, because the difference is which variable is modeled and which is conditioned on. That is decided by the trainer, not by the config:" + "`NLE` uses the same density configs as `NPE`. The trainer determines which variable is modeled and which is conditioned on:" ] }, { @@ -139,7 +139,7 @@ "source": [ "from sbi.neural_nets import NSFConfig\n", "\n", - "NSFConfig(num_bins=12) # fine, a spline flow has bins\n", + "NSFConfig(num_bins=12)\n", "\n", "try:\n", " MAFConfig(num_bins=12)\n", @@ -177,7 +177,7 @@ "source": [ "## Inspecting a config\n", "\n", - "Configs are frozen dataclasses, and their `repr` shows only what you changed, which makes them convenient to log or to put in an experiment record:" + "A config's `repr` shows only non-default fields:" ] }, { @@ -264,7 +264,7 @@ "source": [ "## Settings without a field\n", "\n", - "Some underlying models accept keyword arguments that have no field on the config. Use `extra_kwargs` for these library-specific options. The underlying builder must support the option; this is not a way to make an unsupported setting work:" + "Use `extra_kwargs` for options supported by the underlying library but not exposed as config fields. Keys not listed in the builder signature warn when the estimator is built:" ] }, { @@ -281,7 +281,7 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "A key that duplicates an existing field is rejected, so there is one place a setting can come from:" + "Keys that duplicate config fields are rejected:" ] }, { @@ -343,7 +343,7 @@ "source": [ "## Mixed data for MNPE and MNLE\n", "\n", - "`MNPE` and `MNLE` handle data that is partly continuous and partly discrete. `MixedConfig` describes the whole estimator, and the continuous part is configured by nesting that model's own config, so its settings stay validated by its own class:" + "`MixedConfig` models variables with both continuous and discrete components. Configure its continuous component with a nested density config:" ] }, { @@ -443,7 +443,7 @@ "source": [ "## Building the estimator yourself\n", "\n", - "The trainer calls `build` on the config once it has seen data, because the network needs the shapes and the z-scoring statistics. You can call it yourself if you want the network outside a trainer, for example for a custom training loop as in [the training interface tutorial](../advanced_tutorials/18_training_interface.ipynb):" + "Configs use a data batch to infer shapes and z-scoring statistics. Call `build` directly to use an estimator outside a trainer, for example in a [custom training loop](../advanced_tutorials/18_training_interface.ipynb):" ] }, { From e3dc8b893f4de40269d16cdf59f3b08dd0540fba Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 09:15:53 +0530 Subject: [PATCH 27/29] simplify the density estimator wordings --- .../03_density_estimators.ipynb | 17 ++++------- docs/how_to_guide/03_density_estimators.ipynb | 30 +++---------------- 2 files changed, 9 insertions(+), 38 deletions(-) diff --git a/docs/advanced_tutorials/03_density_estimators.ipynb b/docs/advanced_tutorials/03_density_estimators.ipynb index 433120dac..9ef87cbcf 100644 --- a/docs/advanced_tutorials/03_density_estimators.ipynb +++ b/docs/advanced_tutorials/03_density_estimators.ipynb @@ -27,16 +27,12 @@ "source": [ "## Changing the type of density estimator\n", "\n", - "The density estimator is chosen by passing a **config object** in\n", - "the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`). There is one config class per model, so the choice of model is the choice of class, e.g.\n", - "`MAFConfig` for a Masked Autoregressive Flow, or `NSFConfig` for a Neural Spline Flow with default\n", - "hyperparameters. The config carries only the settings that model accepts, so an unsupported or misspelled one raises at construction instead of being ignored.\n", + "Pass a config object as `density_estimator` to `NPE` or `NLE`, for example `MAFConfig` for a masked autoregressive flow or `NSFConfig` for a neural spline flow.\n", "\n", "Note that `MAFConfig` or `NSFConfig` correspond to `nflows` density\n", "estimators. Those have proven to work well, but the `nflows` package is not maintained\n", "anymore. To use more recent and actively maintained density estimators, we tentatively\n", - "recommend using `zuko`, e.g., `ZukoMAFConfig` or `ZukoNSFConfig`. \n", - "" + "recommend using `zuko`, e.g., `ZukoMAFConfig` or `ZukoNSFConfig`.\n" ] }, { @@ -92,10 +88,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The hyperparameters of a model are the constructor arguments of its config, so you can tune them for the problem at hand.\n", + "Set hyperparameters when constructing the config.\n", "\n", - "Here, because we want to use N*P*E, we pass the config to the `density_estimator` argument of `NPE`. In this example, we will create a neural spline flow (`ZukoNSFConfig`) with `60` hidden units and `3` transform layers:\n", - "" + "Here we configure a Zuko neural spline flow with `60` hidden units and `3` transform layers:\n" ] }, { @@ -104,7 +99,6 @@ "metadata": {}, "outputs": [], "source": [ - "# The same configs are used for NLE. NRE takes classifier configs instead.\n", "from sbi.neural_nets import ZukoNSFConfig\n", "\n", "density_estimator = ZukoNSFConfig(hidden_features=60, num_transforms=3)\n", @@ -117,8 +111,7 @@ "source": [ "It is also possible to pass an `embedding_net` to a config to automatically\n", "learn summary statistics from high-dimensional simulation outputs. You can find a more\n", - "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", - "" + "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n" ] }, { diff --git a/docs/how_to_guide/03_density_estimators.ipynb b/docs/how_to_guide/03_density_estimators.ipynb index 726cf047e..5a43d9d68 100644 --- a/docs/how_to_guide/03_density_estimators.ipynb +++ b/docs/how_to_guide/03_density_estimators.ipynb @@ -27,7 +27,7 @@ "source": [ "## Changing the type of density estimator\n", "\n", - "The density estimator is chosen by passing a **config object** in the `density_estimator` keyword argument to the inference object (`NPE` or `NLE`). There is one config class per model, so the choice of model is the choice of class: `MAFConfig` for a Masked Autoregressive Flow, `NSFConfig` for a Neural Spline Flow, and so on.\n", + "Pass a config object as `density_estimator` to `NPE` or `NLE`, for example `MAFConfig` for a masked autoregressive flow or `NSFConfig` for a neural spline flow.\n", "\n", "Note that `MAFConfig` or `NSFConfig` correspond to `nflows` density\n", "estimators. Those have proven to work well, but the `nflows` package is not maintained\n", @@ -88,9 +88,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "The hyperparameters of a model are the constructor arguments of its config, so you can tune them for the problem at hand.\n", + "Set hyperparameters when constructing the config.\n", "\n", - "Here, because we want to use N*P*E, we pass the config to the `density_estimator` argument of `NPE`. In this example, we will create a neural spline flow (`ZukoNSFConfig`) with `60` hidden units and `3` transform layers:" + "Here we configure a Zuko neural spline flow with `60` hidden units and `3` transform layers:" ] }, { @@ -105,28 +105,6 @@ "inference = NPE(prior=prior, density_estimator=density_estimator)" ] }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "A config only accepts the settings its own model has, so a setting that belongs to a different model, or a misspelled one, raises immediately instead of being ignored:" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "from sbi.neural_nets import ZukoMAFConfig\n", - "\n", - "try:\n", - " # `num_bins` is a spline setting, and a MAF has no splines.\n", - " ZukoMAFConfig(num_bins=8)\n", - "except TypeError as e:\n", - " print(e)" - ] - }, { "cell_type": "markdown", "metadata": {}, @@ -135,7 +113,7 @@ "learn summary statistics from high-dimensional simulation outputs. You can find a more\n", "detailed tutorial on this in [04_embedding_networks](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n", "\n", - "The full list of configs, together with embedding nets, z-scoring, and how to migrate from the older string and factory-function interfaces, is in [how to configure the neural network](27_estimator_configs.ipynb)." + "See the [config guide](27_estimator_configs.ipynb) for available models, field validation, z-scoring, and migration examples." ] }, { From a5269190db364fe9c4d782b675fb988eda88b097 Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 09:16:48 +0530 Subject: [PATCH 28/29] clarify custom builder boundaries --- docs/how_to_guide/24_abstraction_levels.ipynb | 21 ++++++++----------- 1 file changed, 9 insertions(+), 12 deletions(-) diff --git a/docs/how_to_guide/24_abstraction_levels.ipynb b/docs/how_to_guide/24_abstraction_levels.ipynb index 1fcff50ba..97f7765d8 100644 --- a/docs/how_to_guide/24_abstraction_levels.ipynb +++ b/docs/how_to_guide/24_abstraction_levels.ipynb @@ -79,7 +79,7 @@ "\n", "**Use case**: Standard workflows, most common approach\n", "\n", - "Construct a trainer and train. Each method ships a default network that works reasonably across problems, so nothing has to be specified." + "Omit `density_estimator` to use the trainer's default network." ] }, { @@ -121,7 +121,7 @@ "\n", "**Use case**: Choose the model, or tune its architecture hyperparameters\n", "\n", - "Pass a config object as `density_estimator`. There is one class per model, and it carries exactly the settings that model accepts, so an unsupported or misspelled setting is an error at construction rather than something silently dropped." + "Pass a config object as `density_estimator` to select a model and set its hyperparameters. Unsupported or misspelled fields raise an error at construction." ] }, { @@ -133,10 +133,7 @@ "from sbi.neural_nets import MAFConfig\n", "\n", "# Level 2: config object with custom hyperparameters\n", - "density_estimator = MAFConfig(\n", - " hidden_features=50, # Customize hidden layer size\n", - " num_transforms=5, # Customize number of transform layers\n", - ")\n", + "density_estimator = MAFConfig(hidden_features=50, num_transforms=5)\n", "\n", "# Pass to NPE (rest of workflow is the same)\n", "inference = NPE(prior=prior, density_estimator=density_estimator)\n", @@ -161,7 +158,7 @@ "\n", "See [how to configure the neural network](27_estimator_configs.ipynb) for the full list of configs and their settings.\n", "\n", - "The older interfaces for this level still work: a model name as a string, or the `posterior_nn()` / `likelihood_nn()` / `classifier_nn()` factory functions. Strings are deprecated and warn; the migration table is in the same guide." + "Strings and factory functions remain available for compatibility. See the config guide for migration examples." ] }, { @@ -170,7 +167,7 @@ "source": [ "## Level 3: Direct Network Builders\n", "\n", - "**Use case**: Reach a build-function argument that no config exposes\n", + "**Use case**: Use a custom builder or custom estimator\n", "\n", "Pass a callable taking `(batch_theta, batch_x)` and returning an estimator. The example below uses a low-level builder directly for NPE. For NLE, wrap the builder to pass `batch_x` as the modeled input and `batch_theta` as its condition; see the [config guide](27_estimator_configs.ipynb)." ] @@ -252,7 +249,7 @@ "- Must combine with prior using MCMC, VI, or rejection sampling to get posterior samples\n", "- This adds a **second dimension of control**: choosing and configuring the sampling method\n", "\n", - "**Important**: The 4 density estimator levels above work exactly the same for NLE, with the same config classes at Level 2. What changes is which variable is modeled and which is conditioned on, and the trainer decides that, not the config." + "NLE uses the same config classes but models $x$ conditioned on $\\theta$. The trainer handles this role mapping for configs; custom build functions must handle it explicitly." ] }, { @@ -287,9 +284,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "**Note**: Levels 2-4 for the density estimator work identically:\n", + "Levels 2-4 also apply to NLE:\n", "- Level 2: the same config classes, e.g. `NLE(prior, density_estimator=MAFConfig())`\n", - "- Level 3: Use `build_nsf()` (same as NPE)\n", + "- Level 3: Wrap a builder to model `batch_x` conditioned on `batch_theta`\n", "- Level 4: Custom training (see tutorial 18)" ] }, @@ -517,7 +514,7 @@ "| Tune network depth or width | **2** | `ZukoNSFConfig(hidden_features=100)` |\n", "| Add embedding networks for images/timeseries | **2** | `ZukoNSFConfig(embedding_net=my_cnn)` |\n", "| Access specialized flow parameters | **2** | `NSFConfig(num_bins=16, tail_bound=5.0)` |\n", - "| Reach an argument no config exposes | **3** | `partial(build_nsf, ...)` |\n", + "| Use a custom builder or custom estimator | **3** | `partial(build_nsf, ...)` |\n", "| Implement custom network architecture | **3** | Subclass `ConditionalDensityEstimator` |\n", "| Define custom loss functions or training | **4** | See advanced tutorial 18 |\n", "\n", From de0d05ea5539147692210bf2cb37731e949da08d Mon Sep 17 00:00:00 2001 From: Satwik Sai Prakash Sahoo Date: Fri, 18 Sep 2026 09:17:01 +0530 Subject: [PATCH 29/29] fix NRE configuration guidance in the GPU tutorial --- docs/how_to_guide/07_gpu_training.ipynb | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/docs/how_to_guide/07_gpu_training.ipynb b/docs/how_to_guide/07_gpu_training.ipynb index 7d323c971..3b4ad7290 100644 --- a/docs/how_to_guide/07_gpu_training.ipynb +++ b/docs/how_to_guide/07_gpu_training.ipynb @@ -110,13 +110,16 @@ "\n", "Whether or not you reduce your training time when training on a GPU depends on\n", "the problem at hand. We provide a couple of default density estimators for\n", - "`NPE`, `NLE` and `NRE`, e.g., a mixture density network\n", + "`NPE` and `NLE`, e.g., a mixture density network\n", "(`density_estimator=MDNConfig()`) or a Masked Autoregressive Flow\n", "(`density_estimator=MAFConfig()`). For these default density estimators, we do **not**\n", "expect a speed-up. This is because the underlying neural networks are relatively\n", "shallow and not tall, e.g., they do not have many parameters or matrix\n", "operations that benefit from being executed on the GPU.\n", "\n", + "For `NRE`, configure the classifier instead, for example with\n", + "`classifier=ResNetClassifierConfig()`.\n", + "\n", "A speed-up through training on the GPU will most likely become visible when\n", "using convolutional modules in your neural networks. E.g., when passing an\n", "embedding net for image processing like in [this example](https://sbi.readthedocs.io/en/latest/how_to_guide/04_embedding_networks.html).\n",