docs: document the estimator config API

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)."
    ]
   },
   {
@@ -27,15 +27,12 @@
    "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",
+    "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",
-    "**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 +53,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 +72,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 +88,9 @@
    "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",
+    "Set hyperparameters when constructing the config.\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 we configure a Zuko neural spline flow with `60` hidden units and `3` transform layers:\n"
    ]
   },
   {
@@ -98,20 +99,17 @@
    "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",
+    "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"
    ]
@@ -131,7 +129,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",

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)"

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);"
    ]

docs/advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb

@@ -63,6 +63,7 @@
     "from sbi.analysis import pairplot\n",
     "from sbi.inference import NLE, NPE, simulate_for_sbi\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",
@@ -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",
+    ")"
    ]
   },
   {

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()"
    ]

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",

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: