Add RatioEstimatorBuilder for NRE trainers and base class improvements

sbi/inference/trainers/nle/mnle.py

@@ -118,7 +118,8 @@ def __init__(
                 )
             warnings.warn(
                 "Passing a string for `density_estimator` is deprecated. "
-                "Use MixedDensityEstimatorBuilder(...) instead.",
+                "Use MixedDensityEstimatorBuilder(...) instead, e.g. "
+                "`from sbi.neural_nets import MixedDensityEstimatorBuilder`.",
                 FutureWarning,
                 stacklevel=2,
             )

sbi/inference/trainers/nle/nle_base.py

@@ -93,7 +93,8 @@ def __init__(
         elif isinstance(density_estimator, str):
             warnings.warn(
                 "Passing a string for `density_estimator` is deprecated. "
-                "Use DensityEstimatorBuilder(model=...) instead.",
+                "Use DensityEstimatorBuilder(model=...) instead, e.g. "
+                "`from sbi.neural_nets import DensityEstimatorBuilder`.",
                 FutureWarning,
                 stacklevel=3,
             )

sbi/inference/trainers/npe/mnpe.py

@@ -120,7 +120,8 @@ def __init__(
                 )
             warnings.warn(
                 "Passing a string for `density_estimator` is deprecated. "
-                "Use MixedDensityEstimatorBuilder(...) instead.",
+                "Use MixedDensityEstimatorBuilder(...) instead, e.g. "
+                "`from sbi.neural_nets import MixedDensityEstimatorBuilder`.",
                 FutureWarning,
                 stacklevel=2,
             )

sbi/inference/trainers/npe/npe_base.py

@@ -118,7 +118,8 @@ def __init__(
         elif isinstance(density_estimator, str):
             warnings.warn(
                 "Passing a string for `density_estimator` is deprecated. "
-                "Use DensityEstimatorBuilder(model=...) instead.",
+                "Use DensityEstimatorBuilder(model=...) instead, e.g. "
+                "`from sbi.neural_nets import DensityEstimatorBuilder`.",
                 FutureWarning,
                 stacklevel=3,
             )

sbi/inference/trainers/nre/bnre.py

@@ -11,6 +11,7 @@
 from sbi.inference.trainers._contracts import LossArgs, LossArgsBNRE
 from sbi.inference.trainers.nre.nre_a import NRE_A
 from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn
+from sbi.neural_nets.net_builders.estimator_configs import RatioEstimatorBuilder
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import Tracker
 from sbi.utils.sbiutils import del_entries
@@ -65,7 +66,12 @@ class BNRE(NRE_A):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, ConditionalEstimatorBuildFn[RatioEstimator]] = "resnet",
+        classifier: Union[
+            str,
+            RatioEstimatorBuilder,
+            ConditionalEstimatorBuildFn[RatioEstimator],
+            None,
+        ] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -78,13 +84,14 @@ def __init__(
             prior: A probability distribution that expresses prior knowledge about the
                 parameters, e.g. which ranges are meaningful for them. If `None`, the
                 prior must be passed to `.build_posterior()`.
-            classifier: Classifier trained to approximate likelihood ratios. If it is
-                a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet), or a callable that implements the
-                `ConditionalEstimatorBuildFn` protocol. The callable will
-                be called with the first batch of simulations (theta, x), which can thus
-                be used for shape inference and potentially for z-scoring. It returns a
-                `RatioEstimator`.
+            classifier: The classifier used to approximate the
+                likelihood-to-evidence ratio. If ``None`` (default), uses a
+                ``RatioEstimatorBuilder`` with default settings. A
+                ``RatioEstimatorBuilder`` can be passed to configure
+                the classifier. If it is a string (deprecated), use a
+                pre-configured network of the provided type (one of
+                linear, mlp, resnet). Alternatively, a function that
+                builds a custom neural network can be provided.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_a.py

@@ -13,6 +13,7 @@
     RatioEstimatorTrainer,
 )
 from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn
+from sbi.neural_nets.net_builders.estimator_configs import RatioEstimatorBuilder
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import Tracker
 from sbi.utils.sbiutils import del_entries
@@ -63,7 +64,12 @@ class NRE_A(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, ConditionalEstimatorBuildFn[RatioEstimator]] = "resnet",
+        classifier: Union[
+            str,
+            RatioEstimatorBuilder,
+            ConditionalEstimatorBuildFn[RatioEstimator],
+            None,
+        ] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -76,13 +82,14 @@ def __init__(
             prior: A probability distribution that expresses prior knowledge about the
                 parameters, e.g. which ranges are meaningful for them. If `None`, the
                 prior must be passed to `.build_posterior()`.
-            classifier: Classifier trained to approximate likelihood ratios. If it is
-                a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet), or a callable that implements the
-                `ConditionalEstimatorBuildFn` protocol. The callable will
-                be called with the first batch of simulations (theta, x), which can
-                thus be used for shape inference and potentially for z-scoring. It
-                returns a `RatioEstimator`.
+            classifier: The classifier used to approximate the
+                likelihood-to-evidence ratio. If ``None`` (default), uses a
+                ``RatioEstimatorBuilder`` with default settings. A
+                ``RatioEstimatorBuilder`` can be passed to configure
+                the classifier. If it is a string (deprecated), use a
+                pre-configured network of the provided type (one of
+                linear, mlp, resnet). Alternatively, a function that
+                builds a custom neural network can be provided.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_b.py

@@ -13,6 +13,7 @@
     RatioEstimatorTrainer,
 )
 from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn
+from sbi.neural_nets.net_builders.estimator_configs import RatioEstimatorBuilder
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import Tracker
 from sbi.utils.sbiutils import del_entries
@@ -63,7 +64,12 @@ class NRE_B(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, ConditionalEstimatorBuildFn[RatioEstimator]] = "resnet",
+        classifier: Union[
+            str,
+            RatioEstimatorBuilder,
+            ConditionalEstimatorBuildFn[RatioEstimator],
+            None,
+        ] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -76,13 +82,14 @@ def __init__(
             prior: A probability distribution that expresses prior knowledge about the
                 parameters, e.g. which ranges are meaningful for them. If `None`, the
                 prior must be passed to `.build_posterior()`.
-            classifier: Classifier trained to approximate likelihood ratios. If it is
-                a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet), or a callable that implements the
-                `ConditionalEstimatorBuildFn` protocol. The callable will
-                be called with the first batch of simulations (theta, x), which can thus
-                be used for shape inference and potentially for z-scoring. It returns a
-                `RatioEstimator`.
+            classifier: The classifier used to approximate the
+                likelihood-to-evidence ratio. If ``None`` (default), uses a
+                ``RatioEstimatorBuilder`` with default settings. A
+                ``RatioEstimatorBuilder`` can be passed to configure
+                the classifier. If it is a string (deprecated), use a
+                pre-configured network of the provided type (one of
+                linear, mlp, resnet). Alternatively, a function that
+                builds a custom neural network can be provided.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/inference/trainers/nre/nre_base.py

@@ -4,7 +4,7 @@
 import warnings
 from abc import ABC, abstractmethod
 from dataclasses import asdict, replace
-from typing import Any, Dict, Literal, Optional, Sequence, Tuple, Union
+from typing import Any, Callable, Dict, Literal, Optional, Sequence, Tuple, Union
 
 import torch
 from torch import Tensor, eye, ones
@@ -32,6 +32,10 @@
 )
 from sbi.neural_nets import classifier_nn
 from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn
+from sbi.neural_nets.net_builders.estimator_configs import (
+    RatioEstimatorBuilder,
+    _EstimatorBuilderBase,
+)
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import TorchTransform, Tracker
 from sbi.utils import (
@@ -45,7 +49,12 @@ class RatioEstimatorTrainer(NeuralInference[RatioEstimator], ABC):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, ConditionalEstimatorBuildFn[RatioEstimator]] = "resnet",
+        classifier: Union[
+            Literal["linear", "mlp", "resnet"],
+            RatioEstimatorBuilder,
+            ConditionalEstimatorBuildFn[RatioEstimator],
+            None,
+        ] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -70,13 +79,14 @@ def __init__(
           (normalizing constant) of the data $x$.
 
         Args:
-            classifier: Classifier trained to approximate likelihood ratios. If it is
-                a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet), or a callable that implements the
-                `ConditionalEstimatorBuildFn` protocol. The callable will
-                be called with the first batch of simulations (theta, x), which can thus
-                be used for shape inference and potentially for z-scoring. It returns a
-                `RatioEstimator`.
+            classifier: The classifier used to approximate the
+                likelihood-to-evidence ratio. If ``None`` (default), uses a
+                ``RatioEstimatorBuilder`` with default settings. A
+                ``RatioEstimatorBuilder`` can be passed to configure
+                the classifier. If it is a string (deprecated), use a
+                pre-configured network of the provided type (one of
+                linear, mlp, resnet). Alternatively, a function that
+                builds a custom neural network can be provided.
 
         See docstring of `NeuralInference` class for all other arguments.
         """
@@ -90,14 +100,29 @@ def __init__(
             show_progress_bars=show_progress_bars,
         )
 
-        # As detailed in the docstring, `density_estimator` is either a string or
-        # a callable. The function creating the neural network is attached to
-        # `_build_neural_net`. It will be called in the first round and receive
-        # thetas and xs as inputs, so that they can be used for shape inference and
-        # potentially for z-scoring.
-        check_estimator_arg(classifier)
-        if isinstance(classifier, str):
+        if classifier is not None:
+            check_estimator_arg(classifier)
+        if classifier is None:
+            self._build_neural_net = self._wrap_builder(
+                RatioEstimatorBuilder(model="resnet")
+            )
+        elif isinstance(classifier, str):
+            warnings.warn(
+                "Passing a string for `classifier` is deprecated. "
+                "Use RatioEstimatorBuilder(model=...) instead, e.g. "
+                "`from sbi.neural_nets import RatioEstimatorBuilder`.",
+                FutureWarning,
+                stacklevel=3,
+            )
             self._build_neural_net = classifier_nn(model=classifier)
+        elif isinstance(classifier, _EstimatorBuilderBase):
+            if not isinstance(classifier, RatioEstimatorBuilder):
+                raise TypeError(
+                    "NRE requires a RatioEstimatorBuilder; got "
+                    f"{type(classifier).__name__}. Use "
+                    "RatioEstimatorBuilder(model=...)."
+                )
+            self._build_neural_net = self._wrap_builder(classifier)
         else:
             self._build_neural_net = classifier
 
@@ -465,6 +490,22 @@ def _initialize_neural_network(
 
             del x, theta
 
+    @staticmethod
+    def _wrap_builder(
+        builder: _EstimatorBuilderBase,
+    ) -> Callable:
+        """Wrap a builder object as a legacy-compatible build function.
+
+        This allows the existing ``_initialize_neural_network`` flow to work
+        unchanged: the returned callable has the same ``(batch_theta, batch_x)``
+        signature as the functions produced by ``classifier_nn``.
+        """
+
+        def build_fn(batch_theta, batch_x):
+            return builder.build(batch_input=batch_theta, batch_condition=batch_x)
+
+        return build_fn
+
     def _get_losses(self, batch: Sequence[Tensor], loss_args: LossArgs) -> Tensor:
         """
         Compute losses for a batch of data.

sbi/inference/trainers/nre/nre_c.py

@@ -13,6 +13,7 @@
     RatioEstimatorTrainer,
 )
 from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn
+from sbi.neural_nets.net_builders.estimator_configs import RatioEstimatorBuilder
 from sbi.neural_nets.ratio_estimators import RatioEstimator
 from sbi.sbi_types import Tracker
 from sbi.utils.sbiutils import del_entries
@@ -60,7 +61,12 @@ class NRE_C(RatioEstimatorTrainer):
     def __init__(
         self,
         prior: Optional[Distribution] = None,
-        classifier: Union[str, ConditionalEstimatorBuildFn[RatioEstimator]] = "resnet",
+        classifier: Union[
+            str,
+            RatioEstimatorBuilder,
+            ConditionalEstimatorBuildFn[RatioEstimator],
+            None,
+        ] = None,
         device: str = "cpu",
         logging_level: Union[int, str] = "warning",
         summary_writer: Optional[SummaryWriter] = None,
@@ -73,13 +79,14 @@ def __init__(
             prior: A probability distribution that expresses prior knowledge about the
                 parameters, e.g. which ranges are meaningful for them. If `None`, the
                 prior must be passed to `.build_posterior()`.
-            classifier: Classifier trained to approximate likelihood ratios. If it is
-                a string, use a pre-configured network of the provided type (one of
-                linear, mlp, resnet), or a callable that implements the
-                `ConditionalEstimatorBuildFn` protocol. The callable will
-                be called with the first batch of simulations (theta, x), which can thus
-                be used for shape inference and potentially for z-scoring. It returns a
-                `RatioEstimator`.
+            classifier: The classifier used to approximate the
+                likelihood-to-evidence ratio. If ``None`` (default), uses a
+                ``RatioEstimatorBuilder`` with default settings. A
+                ``RatioEstimatorBuilder`` can be passed to configure
+                the classifier. If it is a string (deprecated), use a
+                pre-configured network of the provided type (one of
+                linear, mlp, resnet). Alternatively, a function that
+                builds a custom neural network can be provided.
             device: Training device, e.g., "cpu", "cuda" or "cuda:{0, 1, ...}".
             logging_level: Minimum severity of messages to log. One of the strings
                 INFO, WARNING, DEBUG, ERROR and CRITICAL.

sbi/neural_nets/__init__.py

@@ -9,6 +9,11 @@
     posterior_nn,
     posterior_score_nn,
 )
+from sbi.neural_nets.net_builders.estimator_configs import (
+    DensityEstimatorBuilder,
+    MixedDensityEstimatorBuilder,
+    RatioEstimatorBuilder,
+)
 
 
 def __getattr__(name):
@@ -55,4 +60,7 @@ def __getattr__(name):
     "posterior_nn",
     "posterior_score_nn",
     "posterior_flow_nn",
+    "DensityEstimatorBuilder",
+    "MixedDensityEstimatorBuilder",
+    "RatioEstimatorBuilder",
 ]

sbi/neural_nets/net_builders/__init__.py

@@ -7,7 +7,11 @@
     build_mlp_classifier,
     build_resnet_classifier,
 )
-from sbi.neural_nets.net_builders.estimator_configs import DensityEstimatorBuilder
+from sbi.neural_nets.net_builders.estimator_configs import (
+    DensityEstimatorBuilder,
+    MixedDensityEstimatorBuilder,
+    RatioEstimatorBuilder,
+)
 from sbi.neural_nets.net_builders.flow import (
     build_made,
     build_maf,