Repository: sbi-dev/sbi
Pull request #1362: MNPE class similar to MNLE
Author: login: dgedon
Created: 2025-01-10T15:57:48Z | Merged: 2025-03-20T14:06:24Z

Implementation of mixed NPE where we have some continuous parameters theta followed by one (or multiple with this PR https://github.com/sbi-dev/sbi/pull/1269) discrete parameters. The observation space is fully continuous.

Deprecated mnle.py in net_builders und unified mnle/mnpe as mixed_nets.py.


Diff:
diff --git a/sbi/inference/__init__.py b/sbi/inference/__init__.py
index 7c8b058df..e4abcac01 100644
--- a/sbi/inference/__init__.py
+++ b/sbi/inference/__init__.py
@@ -6,7 +6,7 @@
 )
 from sbi.inference.trainers.fmpe import FMPE
 from sbi.inference.trainers.nle import MNLE, NLE_A
-from sbi.inference.trainers.npe import NPE_A, NPE_B, NPE_C  # noqa: F401
+from sbi.inference.trainers.npe import MNPE, NPE_A, NPE_B, NPE_C  # noqa: F401
 from sbi.inference.trainers.npse import NPSE
 from sbi.inference.trainers.nre import BNRE, NRE_A, NRE_B, NRE_C  # noqa: F401
 
diff --git a/sbi/inference/trainers/npe/__init__.py b/sbi/inference/trainers/npe/__init__.py
index 5d61f8c96..fe13e7774 100644
--- a/sbi/inference/trainers/npe/__init__.py
+++ b/sbi/inference/trainers/npe/__init__.py
@@ -1,3 +1,4 @@
+from sbi.inference.trainers.npe.mnpe import MNPE  # noqa: F401
 from sbi.inference.trainers.npe.npe_a import NPE_A  # noqa: F401
 from sbi.inference.trainers.npe.npe_b import NPE_B  # noqa: F401
 from sbi.inference.trainers.npe.npe_base import PosteriorEstimator  # noqa: F401
diff --git a/sbi/inference/trainers/npe/mnpe.py b/sbi/inference/trainers/npe/mnpe.py
new file mode 100644
index 000000000..8149b8af4
--- /dev/null
+++ b/sbi/inference/trainers/npe/mnpe.py
@@ -0,0 +1,151 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Any, Callable, Dict, Optional, Union
+
+from torch.distributions import Distribution
+
+from sbi.inference.posteriors import (
+    DirectPosterior,
+    ImportanceSamplingPosterior,
+    MCMCPosterior,
+    RejectionPosterior,
+    VIPosterior,
+)
+from sbi.inference.trainers.npe.npe_c import NPE_C
+from sbi.neural_nets.estimators import MixedDensityEstimator
+from sbi.sbi_types import TensorboardSummaryWriter
+from sbi.utils.sbiutils import del_entries
+
+
+class MNPE(NPE_C):
+    def __init__(
+        self,
+        prior: Optional[Distribution] = None,
+        density_estimator: Union[str, Callable] = "mnpe",
+        device: str = "cpu",
+        logging_level: Union[int, str] = "WARNING",
+        summary_writer: Optional[TensorboardSummaryWriter] = None,
+        show_progress_bars: bool = True,
+    ):
+        r"""Mixed Neural Posterior Estimation (MNPE).
+
+        Like NPE-C, but designed to be applied to data with mixed types, e.g.,
+        continuous parameters and discrete parameters like they occur in models with
+        switching components. The emebedding net will only operate on the continuous
+        parameters, note this to design the dimension of the embedding net.
+
+        Args:
+            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()`.
+            density_estimator: If it is a string, it must be "mnpe" to use the
+                preconfigured neural nets for MNPE. Alternatively, a function
+                that builds a custom neural network can be provided. The function 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
+                needs to return a PyTorch `nn.Module` implementing the density
+                estimator. The density estimator needs to provide the methods
+                `.log_prob`, `.log_prob_iid()` and `.sample()`.
+            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.
+            summary_writer: A tensorboard `SummaryWriter` to control, among others, log
+                file location (default is `<current working directory>/logs`.)
+            show_progress_bars: Whether to show a progressbar during simulation and
+                sampling.
+        """
+
+        if isinstance(density_estimator, str):
+            assert (
+                density_estimator == "mnpe"
+            ), f"""MNPE can be used with preconfigured 'mnpe' density estimator only,
+                not with {density_estimator}."""
+        kwargs = del_entries(locals(), entries=("self", "__class__"))
+        super().__init__(**kwargs)
+
+    def train(
+        self,
+        training_batch_size: int = 200,
+        learning_rate: float = 5e-4,
+        validation_fraction: float = 0.1,
+        stop_after_epochs: int = 20,
+        max_num_epochs: int = 2**31 - 1,
+        clip_max_norm: Optional[float] = 5.0,
+        resume_training: bool = False,
+        discard_prior_samples: bool = False,
+        retrain_from_scratch: bool = False,
+        show_train_summary: bool = False,
+        dataloader_kwargs: Optional[Dict] = None,
+    ) -> MixedDensityEstimator:
+        density_estimator = super().train(
+            **del_entries(locals(), entries=("self", "__class__"))
+        )
+        assert isinstance(
+            density_estimator, MixedDensityEstimator
+        ), f"""Internal net must be of type
+            MixedDensityEstimator but is {type(density_estimator)}."""
+        return density_estimator
+
+    def build_posterior(
+        self,
+        density_estimator: Optional[MixedDensityEstimator] = None,
+        prior: Optional[Distribution] = None,
+        sample_with: str = "direct",
+        mcmc_method: str = "slice_np_vectorized",
+        vi_method: str = "rKL",
+        direct_sampling_parameters: Optional[Dict[str, Any]] = None,
+        mcmc_parameters: Optional[Dict[str, Any]] = None,
+        vi_parameters: Optional[Dict[str, Any]] = None,
+        rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
+        importance_sampling_parameters: Optional[Dict[str, Any]] = None,
+    ) -> Union[
+        MCMCPosterior,
+        RejectionPosterior,
+        VIPosterior,
+        DirectPosterior,
+        ImportanceSamplingPosterior,
+    ]:
+        """Build posterior from the neural density estimator.
+
+        Args:
+            density_estimator: The density estimator that the posterior is based on.
+                If `None`, use the latest neural density estimator that was trained.
+            prior: Prior distribution.
+            sample_with: Method to use for sampling from the posterior. Must be one of
+                [`direct` | `mcmc` | `rejection` | `vi` | `importance`].
+            mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
+                `hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
+                implementation of slice sampling; select `hmc`, `nuts` or `slice` for
+                Pyro-based sampling.
+            vi_method: Method used for VI, one of [`rKL`, `fKL`, `IW`, `alpha`].
+            direct_sampling_parameters: Additional kwargs passed to `DirectPosterior`.
+            mcmc_parameters: Additional kwargs passed to `MCMCPosterior`.
+            vi_parameters: Additional kwargs passed to `VIPosterior`.
+            rejection_sampling_parameters: Additional kwargs passed to `
+                RejectionPosterior`.
+            importance_sampling_parameters: Additional kwargs passed to
+                `ImportanceSamplingPosterior`.
+
+        Returns:
+            Posterior $p(\theta|x)$ with `.sample()` and `.log_prob()` methods.
+        """
+        if density_estimator is not None:
+            assert isinstance(
+                density_estimator, MixedDensityEstimator
+            ), f"""net must be of type MixedDensityEstimator but is {
+                type(density_estimator)
+            }."""
+
+        return super().build_posterior(
+            density_estimator=density_estimator,
+            prior=prior,
+            sample_with=sample_with,
+            mcmc_method=mcmc_method,
+            vi_method=vi_method,
+            direct_sampling_parameters=direct_sampling_parameters,
+            mcmc_parameters=mcmc_parameters,
+            vi_parameters=vi_parameters,
+            rejection_sampling_parameters=rejection_sampling_parameters,
+            importance_sampling_parameters=importance_sampling_parameters,
+        )
diff --git a/sbi/neural_nets/estimators/mixed_density_estimator.py b/sbi/neural_nets/estimators/mixed_density_estimator.py
index 0f3ba04d1..b3be6ab81 100644
--- a/sbi/neural_nets/estimators/mixed_density_estimator.py
+++ b/sbi/neural_nets/estimators/mixed_density_estimator.py
@@ -11,10 +11,13 @@
 
 
 class MixedDensityEstimator(ConditionalDensityEstimator):
-    """Class performing Mixed Neural Likelihood Estimation.
+    """Class performing Mixed Neural Density Estimation.
 
-    MNLE combines a Categorical net and a neural density estimator to model data
-    with mixed types, e.g., as they occur in decision-making models.
+    This estimator combines a categorical mass estimator and a density estimator to
+    model variables with mixed types (discrete and continuous). It can be used for both
+    likelihood estimation (e.g., for discrete decisions and continuous reaction times
+    in decision-making models) or posterior estimation (e.g., for models that have both
+    discrete and continuous parameters).
     """
 
     def __init__(
@@ -26,7 +29,8 @@ def __init__(
         embedding_net: nn.Module = nn.Identity(),
         log_transform_input: bool = False,
     ):
-        """Initialize class for combining density estimators for MNLE.
+        """Initialize class for combining density estimators for mixed neural
+        density estimation.
 
         Args:
             discrete_net: neural net to model discrete part of the data.
@@ -51,8 +55,9 @@ def __init__(
 
     def forward(self, input: Tensor):
         raise NotImplementedError(
-            """The forward method is not implemented for MNLE, use '.sample(...)' to
-            generate samples though a forward pass."""
+            """The forward method is not implemented for mixed neural density
+            estimation, use '.sample(...)' to generate samples though a forward
+            pass."""
         )
 
     def sample(
diff --git a/sbi/neural_nets/factory.py b/sbi/neural_nets/factory.py
index e6d65b151..a41ff7b9b 100644
--- a/sbi/neural_nets/factory.py
+++ b/sbi/neural_nets/factory.py
@@ -31,7 +31,7 @@
     build_resnet_flowmatcher,
 )
 from sbi.neural_nets.net_builders.mdn import build_mdn
-from sbi.neural_nets.net_builders.mnle import build_mnle
+from sbi.neural_nets.net_builders.mixed_nets import build_mnle, build_mnpe
 from sbi.neural_nets.net_builders.score_nets import build_score_estimator
 from sbi.utils.nn_utils import check_net_device
 
@@ -42,6 +42,7 @@
     "maf_rqs": build_maf_rqs,
     "nsf": build_nsf,
     "mnle": build_mnle,
+    "mnpe": build_mnpe,
     "zuko_nice": build_zuko_nice,
     "zuko_maf": build_zuko_maf,
     "zuko_nsf": build_zuko_nsf,
diff --git a/sbi/neural_nets/net_builders/__init__.py b/sbi/neural_nets/net_builders/__init__.py
index d5571e936..bd4f5b352 100644
--- a/sbi/neural_nets/net_builders/__init__.py
+++ b/sbi/neural_nets/net_builders/__init__.py
@@ -24,5 +24,5 @@
     build_resnet_flowmatcher,
 )
 from sbi.neural_nets.net_builders.mdn import build_mdn
-from sbi.neural_nets.net_builders.mnle import build_mnle
+from sbi.neural_nets.net_builders.mixed_nets import build_mnle, build_mnpe
 from sbi.neural_nets.net_builders.score_nets import build_score_estimator
diff --git a/sbi/neural_nets/net_builders/mnle.py b/sbi/neural_nets/net_builders/mixed_nets.py
similarity index 74%
rename from sbi/neural_nets/net_builders/mnle.py
rename to sbi/neural_nets/net_builders/mixed_nets.py
index c4bfa224e..60221a6ab 100644
--- a/sbi/neural_nets/net_builders/mnle.py
+++ b/sbi/neural_nets/net_builders/mixed_nets.py
@@ -52,7 +52,7 @@
 }
 
 
-def build_mnle(
+def _build_mixed_density_estimator(
     batch_x: Tensor,
     batch_y: Tensor,
     z_score_x: Optional[str] = "independent",
@@ -69,8 +69,12 @@ def build_mnle(
     tail_bound: float = 10.0,
     log_transform_x: bool = False,
     **kwargs,
-):
-    """Returns a density estimator for mixed data types.
+) -> MixedDensityEstimator:
+    """Base function for building mixed neural density estimators.
+
+    This function contains the shared logic between MNLE and MNPE.
+
+    Returns a density estimator for mixed data types.
 
     Uses an autoregressive categorical density estimator to model the discrete part
     and a conditional density estimator to model the continuous part of the data.
@@ -105,7 +109,7 @@ def build_mnle(
             as z_score_x.
         flow_model: type of flow model to use for the continuous part of the
             data.
-        num_categorical_columns: Number of categorical columns of each variable in the
+        num_categorical_columns: number of categorical columns of each variable in the
             input data. If None, the function will infer this from the data.
         embedding_net: Optional embedding network for y, required if y is > 1D.
         combined_embedding_net: Optional embedding for combining the discrete
@@ -123,18 +127,18 @@ def build_mnle(
         kwargs: additional keyword arguments passed to the flow model.
 
     Returns:
-        MixedDensityEstimator: nn.Module for performing MNLE.
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
     """
-
     check_data_device(batch_x, batch_y)
 
     warnings.warn(
-        "The mixed neural likelihood estimator assumes that x contains "
+        "The mixed neural density estimator assumes that x contains "
         "continuous data in the first n-k columns (e.g., reaction times) and "
         "categorical data in the last k columns (e.g., corresponding choices). If "
         "this is not the case for the passed `x` do not use this function.",
         stacklevel=2,
     )
+
     # Separate continuous and discrete data.
     if num_categorical_columns is None:
         num_disc = int(torch.sum(_is_discrete(batch_x)))
@@ -142,13 +146,16 @@ def build_mnle(
         num_disc = len(num_categorical_columns)
     cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
 
-    # Set up y-embedding net with z-scoring.
+    # The embeedding net is applied to the continuous part of the inputs
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
     if z_score_y_bool:
         embedding_net = nn.Sequential(
             standardizing_net(batch_y, structured_y), embedding_net
         )
-    # Get size of the embedded condition data.
+    else:
+        embedding_net = embedding_net
+
+    # embed
     embedded_batch_y = embedding_net(batch_y)
     combined_condition = torch.cat([disc_x, embedded_batch_y], dim=-1)
 
@@ -156,8 +163,8 @@ def build_mnle(
     discrete_net = build_categoricalmassestimator(
         disc_x,
         batch_y,
-        z_score_x="none",  # discrete data should not be z-scored.
-        z_score_y="none",  # y-embedding net already z-scores.
+        z_score_x="none",  # discrete data should not be z-scored
+        z_score_y="none",  # y-embedding net already z-scores
         num_hidden=hidden_features,
         num_layers=hidden_layers,
         embedding_net=embedding_net,
@@ -176,12 +183,15 @@ def build_mnle(
 
     # Set up a flow for modelling the continuous data, conditioned on the discrete data.
     continuous_net = model_builders[flow_model](
+        # TODO: add support for optional log-transform in flow builders.
         batch_x=(
-            torch.log(cont_x) if log_transform_x else cont_x
-        ),  # log transform manually.
+            torch.log(cont_x + 1e-10)
+            if log_transform_x
+            else cont_x  # can apply log transform if data is strictly positive
+        ),
         batch_y=combined_condition,
         z_score_x=z_score_x,
-        z_score_y="none",  # combined condition is already z-scored.
+        z_score_y="none",  # combined condition is already z-scored
         # combined embedding net for discrete and continuous data.
         embedding_net=combined_embedding_net,
         num_bins=num_bins,
@@ -200,3 +210,53 @@ def build_mnle(
         input_shape=batch_x[0].shape,
         condition_shape=batch_y[0].shape,
     )
+
+
+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    log_transform_x: bool = False,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural likelihood estimator.
+
+    This estimator models p(x|theta) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (data), used to infer dimensionality.
+        batch_y: Batch of ys (parameters), used to infer dimensionality.
+        log_transform_x: whether to apply a log-transform to x. This is by default false
+            because x has to be strictly positive to apply log-transform.
+        **kwargs: Additional arguments passed to _build_mixed_density_estimator.
+
+    Returns:
+        MixedDensityEstimator for MNLE.
+    """
+    return _build_mixed_density_estimator(
+        batch_x=batch_x, batch_y=batch_y, log_transform_x=log_transform_x, **kwargs
+    )
+
+
+def build_mnpe(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    log_transform_x: bool = False,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural posterior estimator.
+
+    This estimator models p(theta|x) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (parameters), used to infer dimensionality.
+        batch_y: Batch of ys (data), used to infer dimensionality.
+        log_transform_x: whether to apply a log-transform to x. This is by default false
+            because x has to be strictly positive to apply log-transform.
+        **kwargs: Additional arguments passed to _build_mixed_density_estimator.
+
+    Returns:
+        MixedDensityEstimator for MNPE.
+    """
+    return _build_mixed_density_estimator(
+        batch_x=batch_x, batch_y=batch_y, log_transform_x=log_transform_x, **kwargs
+    )
diff --git a/sbi/utils/sbiutils.py b/sbi/utils/sbiutils.py
index 3090bc71a..9fbbb1e45 100644
--- a/sbi/utils/sbiutils.py
+++ b/sbi/utils/sbiutils.py
@@ -691,8 +691,25 @@ def mcmc_transform(
     Returns: A transformation that transforms whose `forward()` maps from unconstrained
         (or z-scored) to constrained (or non-z-scored) space.
     """
-
     if enable_transform:
+
+        def prior_mean_std_transform(prior, device):
+            try:
+                prior_mean = prior.mean.to(device)
+                prior_std = prior.stddev.to(device)
+            except (NotImplementedError, AttributeError):
+                warnings.warn(
+                    "The passed discrete prior has no mean or stddev attribute, "
+                    "estimating them from samples to build affine standardizing "
+                    "transform.",
+                    stacklevel=2,
+                )
+                theta = prior.sample(torch.Size((num_prior_samples_for_zscoring,)))
+                prior_mean = theta.mean(dim=0).to(device)
+                prior_std = theta.std(dim=0).to(device)
+            transform = torch_tf.AffineTransform(loc=prior_mean, scale=prior_std)
+            return transform
+
         # Some distributions have a support argument but it raises a
         # NotImplementedError. We catch this case here.
         try:
@@ -718,35 +735,31 @@ def mcmc_transform(
                 constraint = prior.support.base_constraint  # type: ignore
             else:
                 constraint = prior.support
-            if isinstance(constraint, constraints._Real):
-                support_is_bounded = False
+
+            # Check if the support is discrete
+            if hasattr(prior.support, "is_discrete"):
+                is_discrete = prior.support.is_discrete  # type: ignore
             else:
-                support_is_bounded = True
-        else:
-            support_is_bounded = False
+                is_discrete = False
 
-        # Prior with bounded support, e.g., uniform priors.
-        if has_support and support_is_bounded:
-            transform = biject_to(prior.support)
-        # For all other cases build affine transform with mean and std.
+            if is_discrete:
+                # For discrete distributions, use mean/std transform
+                transform = prior_mean_std_transform(prior, device)
+            else:
+                # For continuous distributions, check if support is bounded
+                if isinstance(constraint, constraints._Real):
+                    support_is_bounded = False
+                else:
+                    support_is_bounded = True
+
+                if support_is_bounded:
+                    transform = biject_to(prior.support)
+                else:
+                    # Use mean/std transform for unbounded continuous distributions
+                    transform = prior_mean_std_transform(prior, device)
         else:
-            try:
-                prior_mean = prior.mean.to(device)
-                prior_std = prior.stddev.to(device)
-            except (NotImplementedError, AttributeError):
-                # NotImplementedError -> Distribution that inherits from torch dist but
-                # does not implement mean, e.g., TransformedDistribution.
-                # AttributeError -> Custom distribution that has no mean/std attribute.
-                warnings.warn(
-                    "The passed prior has no mean or stddev attribute, estimating "
-                    "them from samples to build affimed standardizing transform.",
-                    stacklevel=2,
-                )
-                theta = prior.sample(torch.Size((num_prior_samples_for_zscoring,)))
-                prior_mean = theta.mean(dim=0).to(device)
-                prior_std = theta.std(dim=0).to(device)
-
-            transform = torch_tf.AffineTransform(loc=prior_mean, scale=prior_std)
+            # No support property, use mean/std transform
+            transform = prior_mean_std_transform(prior, device)
     else:
         transform = torch_tf.identity_transform
 
diff --git a/sbi/utils/user_input_checks_utils.py b/sbi/utils/user_input_checks_utils.py
index 7c549eed8..722066aea 100644
--- a/sbi/utils/user_input_checks_utils.py
+++ b/sbi/utils/user_input_checks_utils.py
@@ -171,7 +171,11 @@ def __init__(
         dists: Sequence[Distribution],
         validate_args=None,
         arg_constraints: Optional[Dict[str, constraints.Constraint]] = None,
+        device: Optional[str] = None,
     ):
+        if device is not None:
+            raise NotImplementedError("device is not supported yet")
+
         self._check_distributions(dists)
         if validate_args is not None:
             [d.set_default_validate_args(validate_args) for d in dists]
diff --git a/tests/density_estimator_test.py b/tests/density_estimator_test.py
index 9edce954e..6496f11a6 100644
--- a/tests/density_estimator_test.py
+++ b/tests/density_estimator_test.py
@@ -11,9 +11,7 @@
 from torch.distributions import MultivariateNormal
 
 from sbi.neural_nets.embedding_nets import CNNEmbedding
-from sbi.neural_nets.estimators.shape_handling import (
-    reshape_to_sample_batch_event,
-)
+from sbi.neural_nets.estimators.shape_handling import reshape_to_sample_batch_event
 from sbi.neural_nets.net_builders import (
     build_categoricalmassestimator,
     build_made,
@@ -22,6 +20,7 @@
     build_mdn,
     build_mlp_flowmatcher,
     build_mnle,
+    build_mnpe,
     build_nsf,
     build_resnet_flowmatcher,
     build_zuko_bpf,
@@ -316,6 +315,7 @@ def _build_density_estimator_and_tensors(
     # make last dim discrete for mixed density estimators
     batch_input[:, -1] = torch.randint(0, 4, (1000,))
     batch_condition = torch.randn((1000, *condition_event_shape))
+
     if len(condition_event_shape) > 1:
         embedding_net = CNNEmbedding(condition_event_shape, kernel_size=1)
         z_score_y = "structured"
@@ -323,7 +323,11 @@ def _build_density_estimator_and_tensors(
         embedding_net = torch.nn.Identity()
         z_score_y = "independent"
 
-    if density_estimator_build_fn in [build_mnle, build_categoricalmassestimator]:
+    if density_estimator_build_fn in [
+        build_mnle,
+        build_mnpe,
+        build_categoricalmassestimator,
+    ]:
         density_estimator = density_estimator_build_fn(
             batch_x=batch_input,
             batch_y=batch_condition,
@@ -362,6 +366,12 @@ def _build_density_estimator_and_tensors(
         [build_mnle, 1, (2,), (7, 7), 10],
         [build_mnle, 1, (2,), (2, 7), 10],
         [build_mnle, 1, (2,), (7, 2), 10],
+        # Add MNPE test cases (note: x and y roles are swapped)
+        (build_mnpe, 1, (2,), (7,), 1),
+        (build_mnpe, 1, (2,), (7,), 10),
+        [build_mnpe, 1, (2,), (7, 7), 10],
+        [build_mnpe, 1, (2,), (2, 7), 10],
+        [build_mnpe, 1, (2,), (7, 2), 10],
         [build_categoricalmassestimator, 1, (1,), (7, 7), 10],
         [build_categoricalmassestimator, 2, (1,), (7, 7), 10],
         pytest.param(
@@ -374,6 +384,16 @@ def _build_density_estimator_and_tensors(
                 reason="Sample dim > 1 not supported for Mixed Density Estimation"
             ),
         ),
+        pytest.param(
+            build_mnpe,
+            2,
+            (1,),
+            (7,),
+            10,
+            marks=pytest.mark.xfail(
+                reason="Sample dim > 1 not supported for Mixed Density Estimation"
+            ),
+        ),
         pytest.param(
             build_categoricalmassestimator,
             1,
diff --git a/tests/mnpe_test.py b/tests/mnpe_test.py
new file mode 100644
index 000000000..c13acea07
--- /dev/null
+++ b/tests/mnpe_test.py
@@ -0,0 +1,277 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import nn
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_simulator(theta, sigma2=0.1):
+    """Simulator for continuous data with mixed parameters.
+    Returns x~N(mu_z, tau^2)"""
+    device = theta.device
+
+    if isinstance(sigma2, float):
+        sigma2 = torch.tensor(sigma2, device=device)
+
+    mu = theta[:, :2]
+    z = theta[:, -1]
+    # Select mu based on z
+    mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
+    # Sample from N(mu_z, sigma2)
+    x = Normal(loc=mu_z, scale=torch.sqrt(sigma2)).rsample()
+    return x.reshape(-1, 1).to(device)
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize(
+    "device",
+    [
+        pytest.param(
+            "gpu",
+            marks=pytest.mark.xfail(
+                raises=NotImplementedError,
+                reason="MultipleIndependent does not support device handling yet.",
+            ),
+        )
+    ],
+)
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[-0.5, 0.5, 1.0]], device=device)
+    x_o = mixed_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=-torch.ones(2), high=torch.ones(2)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 2),
+            torch.bernoulli(0.8 * torch.ones(batch_size, 1)),
+        ),
+        dim=1,
+    )
+    x_o = mixed_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=-torch.ones(2), high=torch.ones(2)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 3)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_x", ("independent", "none"))
+@pytest.mark.parametrize("embedding_net", (torch.nn.Identity(), FCEmbedding(1, 1)))
+def test_mnpe_api(flow_model: str, z_score_x: str, embedding_net: nn.Module):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[-0.5, 0.5, 1.0]])
+    x_o = mixed_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(-torch.ones(2), torch.ones(2)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_simulator(theta)
+
+    # Build estimator manually
+    # x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_x=z_score_x,
+        embedding_net=embedding_net,
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test different samplers
+    posterior = trainer.build_posterior(prior=prior)
+    posterior.set_default_x(x_o)
+    if isinstance(posterior, VIPosterior):
+        posterior.train().sample((1,))
+    elif isinstance(posterior, RejectionPosterior):
+        posterior.sample((1,))
+    else:
+        posterior.sample((1,))
+
+
+def reference_posterior_mog(tau2, sigma2, mu_0, mu_1, x_o, q, num_posterior_samples):
+    """Reference posterior for a mixture of two Gaussians.
+    Args:
+        tau2: prior variance of mu0, mu1
+        sigma2: variance of x
+        mu_0: prior mean of mu0
+        mu_1: prior mean of mu1
+        x_o: observation
+        q: probability of z=1
+        num_posterior_samples: number of samples from the posterior
+    """
+    prior_mu0 = Normal(loc=mu_0, scale=torch.sqrt(tau2))
+    prior_mu1 = Normal(loc=mu_1, scale=torch.sqrt(tau2))
+
+    s2 = tau2 + sigma2
+
+    ###### Posterior of z ######
+    normal_0 = Normal(loc=mu_0, scale=torch.sqrt(s2))
+    normal_1 = Normal(loc=mu_1, scale=torch.sqrt(s2))
+    likelihood_0 = normal_0.log_prob(x_o).exp()  # p(x|z=0)
+    likelihood_1 = normal_1.log_prob(x_o).exp()  # p(x|z=1)
+    numerator = q * likelihood_1
+    denominator = numerator + (1 - q) * likelihood_0
+    bernoulli_probs = numerator / denominator
+    posterior_z_ref = Bernoulli(probs=bernoulli_probs)
+    posterior_z_ref_samples = posterior_z_ref.sample((num_posterior_samples,))
+    posterior_z_ref_samples = posterior_z_ref_samples.reshape(-1, 1)
+
+    ###### Posterior of mu ######
+    # posterior std
+    tau2_ref = (1 / tau2 + 1 / sigma2).pow(-1)
+    # posterior mean
+    mu_0_ref = tau2_ref * (mu_0 / tau2 + x_o / sigma2)
+    mu_1_ref = tau2_ref * (mu_1 / tau2 + x_o / sigma2)
+    # posterior distribution
+    posterior_mu0_ref = Normal(loc=mu_0_ref, scale=torch.sqrt(tau2_ref))
+    posterior_mu1_ref = Normal(loc=mu_1_ref, scale=torch.sqrt(tau2_ref))
+
+    # Create binary mask from z samples (True where z=0, False where z=1)
+    z_mask = (posterior_z_ref_samples == 0).flatten()
+
+    # Create empty tensor for results
+    posterior_ref_samples = torch.zeros(num_posterior_samples, 3)
+
+    # Sample and populate directly without broadcasting issues
+    posterior_ref_samples[:, 2] = posterior_z_ref_samples.flatten()
+
+    # For mu0 column: use posterior_mu0_ref when z=0, prior_mu0 when z=1
+    posterior_ref_samples[z_mask, 0] = posterior_mu0_ref.sample((
+        z_mask.sum(),
+    )).flatten()
+    posterior_ref_samples[~z_mask, 0] = prior_mu0.sample(((~z_mask).sum(),)).flatten()
+
+    # For mu1 column: use prior_mu1 when z=0, posterior_mu1_ref when z=1
+    posterior_ref_samples[z_mask, 1] = prior_mu1.sample((z_mask.sum(),)).flatten()
+    posterior_ref_samples[~z_mask, 1] = posterior_mu1_ref.sample((
+        (~z_mask).sum(),
+    )).flatten()
+
+    return posterior_ref_samples
+
+
+@pytest.mark.slow
+def test_mnpe_accuracy():
+    """Test MNPE accuracy."""
+    num_simulations = 1000
+    num_posterior_samples = 1000
+
+    # Prior
+    mu_0 = torch.tensor([-1.0])
+    mu_1 = torch.tensor([1.0])
+    tau2 = torch.tensor([0.1])
+    prior_mu0 = Normal(loc=mu_0, scale=torch.sqrt(tau2))
+    prior_mu1 = Normal(loc=mu_1, scale=torch.sqrt(tau2))
+
+    q = torch.tensor(0.8)
+    prior_z = Bernoulli(probs=q * torch.ones(1))
+
+    prior = MultipleIndependent(
+        [prior_mu0, prior_mu1, prior_z],
+    )
+
+    # Simulator
+    sigma2 = torch.tensor(0.1)
+
+    # observation
+    theta_o = prior.sample((1,))
+    x_o = mixed_simulator(theta_o, sigma2)
+
+    # Reference posterior
+    posterior_ref_samples = reference_posterior_mog(
+        tau2, sigma2, mu_0, mu_1, x_o, q, num_posterior_samples
+    )
+
+    # MNPE
+    theta = prior.sample((num_simulations,))
+    x = mixed_simulator(theta, sigma2)
+
+    trainer = MNPE()
+    trainer.append_simulations(theta, x).train()
+
+    posterior_mnpe = trainer.build_posterior(prior=prior)
+    posterior_mnpe.set_default_x(x_o)
+
+    posterior_mnpe_samples = posterior_mnpe.sample((num_posterior_samples,))
+
+    # Check C2ST
+    check_c2st(
+        posterior_mnpe_samples,
+        posterior_ref_samples,
+        alg="MNPE",
+    )


Reviews:
user: login: janfb
body: looks good overall, except one central question about the call signature of `MNPE`. 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-18T22:20:49Z
user: login: dgedon
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T07:29:44Z
user: login: dgedon
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T08:36:49Z
user: login: janfb
body: some more comments on the tests and the refactoring. 

I am suggestion a toy example with ground-truth posterior for the MNPE scenario to test the accuracy. 
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-19T09:16:02Z
user: login: dgedon
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-19T10:38:42Z
user: login: janfb
body: great to have the test now! 🎉 

added a some comments for the mnle defaults and some suggestions for the test.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T10:23:40Z
user: login: dgedon
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T11:01:16Z
user: login: janfb
body: looks good! 

two final comments.
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T11:26:11Z
user: login: janfb
body: Two final final comments. Good to be merged afterwards. Really cool to have this implemented now 🔥 🎉 
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-03-20T11:49:43Z

Review comments:
diff_hunk: @@ -11,10 +11,12 @@
 
 
 class MixedDensityEstimator(ConditionalDensityEstimator):
-    """Class performing Mixed Neural Likelihood Estimation.
+    """Class performing Mixed Neural Density Estimation.
path: sbi/neural_nets/estimators/mixed_density_estimator.py
user: login: janfb
body: 👍 
created_at: 2025-03-18T13:54:44Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,151 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Any, Callable, Dict, Optional, Union
+
+from torch.distributions import Distribution
+
+from sbi.inference.posteriors import (
+    DirectPosterior,
+    ImportanceSamplingPosterior,
+    MCMCPosterior,
+    RejectionPosterior,
+    VIPosterior,
+)
+from sbi.inference.trainers.npe.npe_c import NPE_C
+from sbi.neural_nets.estimators import MixedDensityEstimator
+from sbi.sbi_types import TensorboardSummaryWriter, TorchModule
+from sbi.utils.sbiutils import del_entries
+
+
+class MNPE(NPE_C):
+    def __init__(
+        self,
+        prior: Optional[Distribution] = None,
+        density_estimator: Union[str, Callable] = "mnpe",
+        device: str = "cpu",
+        logging_level: Union[int, str] = "WARNING",
+        summary_writer: Optional[TensorboardSummaryWriter] = None,
+        show_progress_bars: bool = True,
+    ):
+        r"""Mixed Neural Posterior Estimation (MNPE).
+
+        Like NPE-C, but designed to be applied to data with mixed types, e.g.,
+        continuous parameters and discrete parameters like they occur in models with
+        switching components. The emebedding net will only operate on the continuous
+        parameters, note this to design the dimension of the embedding net.
+
+        Args:
+            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()`.
+            density_estimator: If it is a string, it must be "mnpe" to use the
+                preconfigured neural nets for MNPE. Alternatively, a function
+                that builds a custom neural network can be provided. The function 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
+                needs to return a PyTorch `nn.Module` implementing the density
+                estimator. The density estimator needs to provide the methods
+                `.log_prob`, `.log_prob_iid()` and `.sample()`.
+            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.
+            summary_writer: A tensorboard `SummaryWriter` to control, among others, log
+                file location (default is `<current working directory>/logs`.)
+            show_progress_bars: Whether to show a progressbar during simulation and
+                sampling.
+        """
+
+        if isinstance(density_estimator, str):
+            assert (
+                density_estimator == "mnpe"
+            ), f"""MNPE can be used with preconfigured 'mnpe' density estimator only,
+                not with {density_estimator}."""
+        kwargs = del_entries(locals(), entries=("self", "__class__"))
+        super().__init__(**kwargs)
+
+    def train(
+        self,
+        training_batch_size: int = 200,
+        learning_rate: float = 5e-4,
+        validation_fraction: float = 0.1,
+        stop_after_epochs: int = 20,
+        max_num_epochs: int = 2**31 - 1,
+        clip_max_norm: Optional[float] = 5.0,
+        resume_training: bool = False,
+        discard_prior_samples: bool = False,
+        retrain_from_scratch: bool = False,
+        show_train_summary: bool = False,
+        dataloader_kwargs: Optional[Dict] = None,
+    ) -> MixedDensityEstimator:
+        density_estimator = super().train(
+            **del_entries(locals(), entries=("self", "__class__"))
+        )
+        assert isinstance(
+            density_estimator, MixedDensityEstimator
+        ), f"""Internal net must be of type
+            MixedDensityEstimator but is {type(density_estimator)}."""
+        return density_estimator
+
+    def build_posterior(
+        self,
+        density_estimator: Optional[TorchModule] = None,
+        prior: Optional[Distribution] = None,
+        sample_with: str = "direct",
+        mcmc_method: str = "slice_np_vectorized",
+        vi_method: str = "rKL",
+        direct_sampling_parameters: Optional[Dict[str, Any]] = None,
+        mcmc_parameters: Optional[Dict[str, Any]] = None,
+        vi_parameters: Optional[Dict[str, Any]] = None,
+        rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
+        importance_sampling_parameters: Optional[Dict[str, Any]] = None,
+    ) -> Union[
+        MCMCPosterior,
+        RejectionPosterior,
+        VIPosterior,
+        DirectPosterior,
+        ImportanceSamplingPosterior,
+    ]:
+        """Build posterior from the neural density estimator.
+
+        Args:
+            density_estimator: The density estimator that the posterior is based on.
+                If `None`, use the latest neural density estimator that was trained.
+            prior: Prior distribution.
+            sample_with: Method to use for sampling from the posterior. Must be one of
+                [`direct` | `mcmc` | `rejection` | `vi` | `importance`].
+            mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
+                `hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
+                implementation of slice sampling; select `hmc`, `nuts` or `slice` for
+                Pyro-based sampling.
+            vi_method: Method used for VI, one of [`rKL`, `fKL`, `IW`, `alpha`].
+            direct_sampling_parameters: Additional kwargs passed to `DirectPosterior`.
+            mcmc_parameters: Additional kwargs passed to `MCMCPosterior`.
+            vi_parameters: Additional kwargs passed to `VIPosterior`.
+            rejection_sampling_parameters: Additional kwargs passed to `
+                RejectionPosterior`.
+            importance_sampling_parameters: Additional kwargs passed to
+                `ImportanceSamplingPosterior`.
+
+        Returns:
+            Posterior $p(\theta|x)$ with `.sample()` and `.log_prob()` methods.
+        """
+        if density_estimator is not None:
+            assert isinstance(
+                density_estimator, MixedDensityEstimator
+            ), f"""net must be of type MixedDensityEstimator but is {
+                type(density_estimator)
+            }."""
path: sbi/inference/trainers/npe/mnpe.py
user: login: janfb
body: we could also change the type above to be `MixedDensityEstimator` to have a static check. 
created_at: 2025-03-18T13:55:30Z
author_association: CONTRIBUTOR
diff_hunk: @@ -11,10 +11,12 @@
 
 
 class MixedDensityEstimator(ConditionalDensityEstimator):
-    """Class performing Mixed Neural Likelihood Estimation.
+    """Class performing Mixed Neural Density Estimation.
 
-    MNLE combines a Categorical net and a neural density estimator to model data
-    with mixed types, e.g., as they occur in decision-making models.
+    This estimator combines a Categorical net and a neural density estimator to model
+    data with mixed types (discrete and continuous), e.g., as they occur in
+    decision-making models. It can be used for both likelihood and posterior estimation
+    of mixed data.
path: sbi/neural_nets/estimators/mixed_density_estimator.py
user: login: janfb
body: ```suggestion
    This estimator combines a categorical mass estimator and a density estimator to model
    variables with mixed types (discrete and continuous). It can be used for both likelihood
    estimation (e.g., for discrete decisions and continuous reaction times in decision-making
    models) or posterior estimation (e.g., for models that have both discrete and continuous
    parameters). 
```
created_at: 2025-03-18T13:59:21Z
author_association: CONTRIBUTOR
diff_hunk: @@ -51,8 +54,9 @@ def __init__(
 
     def forward(self, input: Tensor):
         raise NotImplementedError(
-            """The forward method is not implemented for MNLE, use '.sample(...)' to
-            generate samples though a forward pass."""
+            """The forward method is not implemented for mixed neural density
+            estimation,use '.sample(...)' to generate samples though a forward
path: sbi/neural_nets/estimators/mixed_density_estimator.py
user: login: janfb
body: ```suggestion
            estimation, use '.sample(...)' to generate samples though a forward
```
created_at: 2025-03-18T13:59:49Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,269 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Literal, Optional
+
+import torch
+from torch import Tensor, nn
+
+from sbi.neural_nets.estimators import MixedDensityEstimator
+from sbi.neural_nets.estimators.mixed_density_estimator import (
+    _is_discrete,
+    _separate_input,
+)
+from sbi.neural_nets.net_builders.categorial import (
+    build_categoricalmassestimator,
+)
+from sbi.neural_nets.net_builders.flow import (
+    build_made,
+    build_maf,
+    build_maf_rqs,
+    build_nsf,
+    build_zuko_bpf,
+    build_zuko_gf,
+    build_zuko_maf,
+    build_zuko_naf,
+    build_zuko_ncsf,
+    build_zuko_nice,
+    build_zuko_nsf,
+    build_zuko_sospf,
+    build_zuko_unaf,
+)
+from sbi.neural_nets.net_builders.mdn import build_mdn
+from sbi.utils.sbiutils import standardizing_net, z_score_parser
+from sbi.utils.user_input_checks import check_data_device
+
+model_builders = {
+    "mdn": build_mdn,
+    "made": build_made,
+    "maf": build_maf,
+    "maf_rqs": build_maf_rqs,
+    "nsf": build_nsf,
+    "zuko_nice": build_zuko_nice,
+    "zuko_maf": build_zuko_maf,
+    "zuko_nsf": build_zuko_nsf,
+    "zuko_ncsf": build_zuko_ncsf,
+    "zuko_sospf": build_zuko_sospf,
+    "zuko_naf": build_zuko_naf,
+    "zuko_unaf": build_zuko_unaf,
+    "zuko_gf": build_zuko_gf,
+    "zuko_bpf": build_zuko_bpf,
+}
+
+
+def _build_mixed_density_estimator(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    z_score_x: Optional[str] = "independent",
+    z_score_y: Optional[str] = "independent",
+    flow_model: str = "nsf",
+    num_categorical_columns: Optional[Tensor] = None,
+    embedding_net: nn.Module = nn.Identity(),
+    combined_embedding_net: Optional[nn.Module] = None,
+    num_transforms: int = 2,
+    num_components: int = 5,
+    num_bins: int = 5,
+    hidden_features: int = 50,
+    hidden_layers: int = 2,
+    tail_bound: float = 10.0,
+    log_transform_x: bool = True,
+    mode: Literal["mnle", "mnpe"] = "mnle",
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Base function for building mixed neural density estimators.
+
+    This function contains the shared logic between MNLE and MNPE.
+
+    Returns a density estimator for mixed data types.
+
+    Uses an autoregressive categorical density estimator to model the discrete part
+    and a conditional density estimator to model the continuous part of the data.
+
+    Note: If the condition y is > 1D, an embedding net must be provided. Then,
+    during inference, we need to combine the embedded condition with the
+    discrete part of the input data. To do this, we use a combined embedding net
+    that takes the discrete part of the input and the embedded condition as
+    input and passes it on to the flow model.
+    To this end, we build the z-scoring and the embedding net in this function
+    here, i.e., outside of the flow model, and then pass z_score_x="none" to the
+    flow builder.
+    The y-embedding is passed to MixedDensityEstimator, which then uses it to
+    embed the continuous part and concatenate it with the discrete part during
+    log_prob evaluation and sampling.
+    The combined embedding net is passed to the flow model, which then uses it
+    to combine the already embedded and combined condition.
+
+    Args:
+        batch_x: Batch of xs, used to infer dimensionality.
+        batch_y: Batch of ys, used to infer dimensionality and (optional)
+            z-scoring.
+        z_score_x: Whether to z-score xs passing into the network, can be one of:
+            - `none`, or None: do not z-score.
+            - `independent`: z-score each dimension independently.
+            - `structured`: treat dimensions as related, therefore compute mean
+              and std
+            over the entire batch, instead of per-dimension. Should be used when
+            each sample is, for example, a time series or an image. For MNLE,
+            this applies to the continuous part of the data only.
+        z_score_y: Whether to z-score ys passing into the network, same options
+            as z_score_x.
+        flow_model: type of flow model to use for the continuous part of the
+            data.
+        num_categorical_columns: number of categorical columns of each variable in the
+            input data. If None, the function will infer this from the data.
+        embedding_net: Optional embedding network for y, required if y is > 1D.
+        combined_embedding_net: Optional embedding for combining the discrete
+            part of the input and the embedded condition into a joined
+            condition. If None, a shallow MLP is used.
+        num_transforms: number of transforms in the flow model.
+        num_components: number of components in the mixture model.
+        num_bins: bins per spline for NSF.
+        hidden_features: number of hidden features used in both nets.
+        hidden_layers: number of hidden layers in the categorical net.
+        tail_bound: spline tail bound for NSF.
+        log_transform_x: whether to apply a log-transform to x to move it to unbounded
+            space, e.g., in case x consists of reaction time data (bounded by
+            zero).
+        kwargs: additional keyword arguments passed to the flow model.
+
+    Returns:
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
+    """
+    check_data_device(batch_x, batch_y)
+
+    warnings.warn(
+        "The mixed neural density estimator assumes that x contains "
+        "continuous data in the first n-k columns (e.g., reaction times) and "
+        "categorical data in the last k columns (e.g., corresponding choices). If "
+        "this is not the case for the passed `x` do not use this function.",
+        stacklevel=2,
+    )
+
+    # Separate continuous and discrete data.
+    if num_categorical_columns is None:
+        num_disc = int(torch.sum(_is_discrete(batch_x)))
+    else:
+        num_disc = len(num_categorical_columns)
+    cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
+
+    # The embedding net is applied to the continuous part of the parameters
+    # MNPE: parameters are in batch_x, MNLE: parameters are in batch_y
+    if mode.lower() == "mnpe":
+        z_score_x_bool, structured_x = z_score_parser(z_score_x)
+        if z_score_x_bool:
+            embedding_net_x = nn.Sequential(
+                standardizing_net(cont_x, structured_x), embedding_net
+            )
+        else:
+            embedding_net_x = embedding_net
+        embedding_net_y = standardizing_net(batch_y)
+    elif mode.lower() == "mnle":
+        z_score_y_bool, structured_y = z_score_parser(z_score_y)
+        if z_score_y_bool:
+            embedding_net_y = nn.Sequential(
+                standardizing_net(batch_y, structured_y), embedding_net
+            )
+        else:
+            embedding_net_y = embedding_net
+        embedding_net_x = nn.Identity()
+    else:
+        raise ValueError(f"Invalid mode: {mode}, must be one of 'mnle' or 'mnpe'.")
+
+    # embed
+    embedded_batch_x = embedding_net_x(cont_x)
+    embedded_batch_y = embedding_net_y(batch_y)
+    combined_condition = torch.cat([disc_x, embedded_batch_y], dim=-1)
+
+    # Set up a categorical RV neural net for modelling the discrete data.
+    discrete_net = build_categoricalmassestimator(
+        disc_x,
+        batch_y,
+        z_score_x="none",
+        z_score_y="none",
+        num_hidden=hidden_features,
+        num_layers=hidden_layers,
+        embedding_net=embedding_net,
+        num_categories=num_categorical_columns,
+    )
+
+    if combined_embedding_net is None:
+        # set up linear embedding net for combining discrete and continuous
+        # data.
+        combined_embedding_net = nn.Sequential(
+            nn.Linear(combined_condition.shape[-1], hidden_features),
+            nn.ReLU(),
+            nn.Linear(hidden_features, hidden_features),
+            nn.ReLU(),
+        )
+
+    # Set up a flow for modelling the continuous data, conditioned on the discrete data.
+    continuous_net = model_builders[flow_model](
+        # TODO: add support for optional log-transform in flow builders.
+        batch_x=(
+            torch.log(embedded_batch_x + 1e-10) if log_transform_x else embedded_batch_x
+        ),  # log transform manually.
+        batch_y=combined_condition,
+        z_score_x="none" if mode.lower() == "mnpe" else z_score_x,
+        z_score_y="none",  # combined condition is already z-scored for mnle
+        # combined embedding net for discrete and continuous data.
+        embedding_net=combined_embedding_net,
+        num_bins=num_bins,
+        num_transforms=num_transforms,
+        num_components=num_components,
+        tail_bound=tail_bound,
+        hidden_features=hidden_features,
+        **kwargs,
+    )
+
+    return MixedDensityEstimator(
+        discrete_net=discrete_net,
+        continuous_net=continuous_net,
+        embedding_net=embedding_net_y,  # pass embedding for continuous condition part.
+        log_transform_input=log_transform_x,
+        input_shape=batch_x[0].shape,
+        condition_shape=batch_y[0].shape,
+    )
+
+
+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural likelihood estimator.
+
+    This estimator models p(x|theta) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (data), used to infer dimensionality.
+        batch_y: Batch of ys (parameters), used to infer dimensionality.
+        **kwargs: Additional arguments passed to _build_mixed_density_estimator.
+
+    Returns:
+        MixedDensityEstimator for MNLE.
+    """
+    return _build_mixed_density_estimator(
+        batch_x=batch_x, batch_y=batch_y, mode="mnle", **kwargs
+    )
+
+
+def build_mnpe(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural posterior estimator.
+
+    This estimator models p(theta|x) where x contains both continuous and discrete data.
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: I am confused by these two call functions. maybe I am missing something, but couldn't we just call `_build_mixed_density_estimator` with `batch_x` and `batch_y` swapped for `MNPE` and `MNLE`?

To me it seems this swapping is not happening, i.e., we need to make sure that in `MNPE` we are only embedding `x` and not `theta`, and v.v. in `MNLE`. Let's discuss tomorrow. 
created_at: 2025-03-18T22:20:20Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,269 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Literal, Optional
+
+import torch
+from torch import Tensor, nn
+
+from sbi.neural_nets.estimators import MixedDensityEstimator
+from sbi.neural_nets.estimators.mixed_density_estimator import (
+    _is_discrete,
+    _separate_input,
+)
+from sbi.neural_nets.net_builders.categorial import (
+    build_categoricalmassestimator,
+)
+from sbi.neural_nets.net_builders.flow import (
+    build_made,
+    build_maf,
+    build_maf_rqs,
+    build_nsf,
+    build_zuko_bpf,
+    build_zuko_gf,
+    build_zuko_maf,
+    build_zuko_naf,
+    build_zuko_ncsf,
+    build_zuko_nice,
+    build_zuko_nsf,
+    build_zuko_sospf,
+    build_zuko_unaf,
+)
+from sbi.neural_nets.net_builders.mdn import build_mdn
+from sbi.utils.sbiutils import standardizing_net, z_score_parser
+from sbi.utils.user_input_checks import check_data_device
+
+model_builders = {
+    "mdn": build_mdn,
+    "made": build_made,
+    "maf": build_maf,
+    "maf_rqs": build_maf_rqs,
+    "nsf": build_nsf,
+    "zuko_nice": build_zuko_nice,
+    "zuko_maf": build_zuko_maf,
+    "zuko_nsf": build_zuko_nsf,
+    "zuko_ncsf": build_zuko_ncsf,
+    "zuko_sospf": build_zuko_sospf,
+    "zuko_naf": build_zuko_naf,
+    "zuko_unaf": build_zuko_unaf,
+    "zuko_gf": build_zuko_gf,
+    "zuko_bpf": build_zuko_bpf,
+}
+
+
+def _build_mixed_density_estimator(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    z_score_x: Optional[str] = "independent",
+    z_score_y: Optional[str] = "independent",
+    flow_model: str = "nsf",
+    num_categorical_columns: Optional[Tensor] = None,
+    embedding_net: nn.Module = nn.Identity(),
+    combined_embedding_net: Optional[nn.Module] = None,
+    num_transforms: int = 2,
+    num_components: int = 5,
+    num_bins: int = 5,
+    hidden_features: int = 50,
+    hidden_layers: int = 2,
+    tail_bound: float = 10.0,
+    log_transform_x: bool = True,
+    mode: Literal["mnle", "mnpe"] = "mnle",
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Base function for building mixed neural density estimators.
+
+    This function contains the shared logic between MNLE and MNPE.
+
+    Returns a density estimator for mixed data types.
+
+    Uses an autoregressive categorical density estimator to model the discrete part
+    and a conditional density estimator to model the continuous part of the data.
+
+    Note: If the condition y is > 1D, an embedding net must be provided. Then,
+    during inference, we need to combine the embedded condition with the
+    discrete part of the input data. To do this, we use a combined embedding net
+    that takes the discrete part of the input and the embedded condition as
+    input and passes it on to the flow model.
+    To this end, we build the z-scoring and the embedding net in this function
+    here, i.e., outside of the flow model, and then pass z_score_x="none" to the
+    flow builder.
+    The y-embedding is passed to MixedDensityEstimator, which then uses it to
+    embed the continuous part and concatenate it with the discrete part during
+    log_prob evaluation and sampling.
+    The combined embedding net is passed to the flow model, which then uses it
+    to combine the already embedded and combined condition.
+
+    Args:
+        batch_x: Batch of xs, used to infer dimensionality.
+        batch_y: Batch of ys, used to infer dimensionality and (optional)
+            z-scoring.
+        z_score_x: Whether to z-score xs passing into the network, can be one of:
+            - `none`, or None: do not z-score.
+            - `independent`: z-score each dimension independently.
+            - `structured`: treat dimensions as related, therefore compute mean
+              and std
+            over the entire batch, instead of per-dimension. Should be used when
+            each sample is, for example, a time series or an image. For MNLE,
+            this applies to the continuous part of the data only.
+        z_score_y: Whether to z-score ys passing into the network, same options
+            as z_score_x.
+        flow_model: type of flow model to use for the continuous part of the
+            data.
+        num_categorical_columns: number of categorical columns of each variable in the
+            input data. If None, the function will infer this from the data.
+        embedding_net: Optional embedding network for y, required if y is > 1D.
+        combined_embedding_net: Optional embedding for combining the discrete
+            part of the input and the embedded condition into a joined
+            condition. If None, a shallow MLP is used.
+        num_transforms: number of transforms in the flow model.
+        num_components: number of components in the mixture model.
+        num_bins: bins per spline for NSF.
+        hidden_features: number of hidden features used in both nets.
+        hidden_layers: number of hidden layers in the categorical net.
+        tail_bound: spline tail bound for NSF.
+        log_transform_x: whether to apply a log-transform to x to move it to unbounded
+            space, e.g., in case x consists of reaction time data (bounded by
+            zero).
+        kwargs: additional keyword arguments passed to the flow model.
+
+    Returns:
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
+    """
+    check_data_device(batch_x, batch_y)
+
+    warnings.warn(
+        "The mixed neural density estimator assumes that x contains "
+        "continuous data in the first n-k columns (e.g., reaction times) and "
+        "categorical data in the last k columns (e.g., corresponding choices). If "
+        "this is not the case for the passed `x` do not use this function.",
+        stacklevel=2,
+    )
+
+    # Separate continuous and discrete data.
+    if num_categorical_columns is None:
+        num_disc = int(torch.sum(_is_discrete(batch_x)))
+    else:
+        num_disc = len(num_categorical_columns)
+    cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
+
+    # The embedding net is applied to the continuous part of the parameters
+    # MNPE: parameters are in batch_x, MNLE: parameters are in batch_y
+    if mode.lower() == "mnpe":
+        z_score_x_bool, structured_x = z_score_parser(z_score_x)
+        if z_score_x_bool:
+            embedding_net_x = nn.Sequential(
+                standardizing_net(cont_x, structured_x), embedding_net
+            )
+        else:
+            embedding_net_x = embedding_net
+        embedding_net_y = standardizing_net(batch_y)
+    elif mode.lower() == "mnle":
+        z_score_y_bool, structured_y = z_score_parser(z_score_y)
+        if z_score_y_bool:
+            embedding_net_y = nn.Sequential(
+                standardizing_net(batch_y, structured_y), embedding_net
+            )
+        else:
+            embedding_net_y = embedding_net
+        embedding_net_x = nn.Identity()
+    else:
+        raise ValueError(f"Invalid mode: {mode}, must be one of 'mnle' or 'mnpe'.")
+
+    # embed
+    embedded_batch_x = embedding_net_x(cont_x)
+    embedded_batch_y = embedding_net_y(batch_y)
+    combined_condition = torch.cat([disc_x, embedded_batch_y], dim=-1)
+
+    # Set up a categorical RV neural net for modelling the discrete data.
+    discrete_net = build_categoricalmassestimator(
+        disc_x,
+        batch_y,
+        z_score_x="none",
+        z_score_y="none",
+        num_hidden=hidden_features,
+        num_layers=hidden_layers,
+        embedding_net=embedding_net,
+        num_categories=num_categorical_columns,
+    )
+
+    if combined_embedding_net is None:
+        # set up linear embedding net for combining discrete and continuous
+        # data.
+        combined_embedding_net = nn.Sequential(
+            nn.Linear(combined_condition.shape[-1], hidden_features),
+            nn.ReLU(),
+            nn.Linear(hidden_features, hidden_features),
+            nn.ReLU(),
+        )
+
+    # Set up a flow for modelling the continuous data, conditioned on the discrete data.
+    continuous_net = model_builders[flow_model](
+        # TODO: add support for optional log-transform in flow builders.
+        batch_x=(
+            torch.log(embedded_batch_x + 1e-10) if log_transform_x else embedded_batch_x
+        ),  # log transform manually.
+        batch_y=combined_condition,
+        z_score_x="none" if mode.lower() == "mnpe" else z_score_x,
+        z_score_y="none",  # combined condition is already z-scored for mnle
+        # combined embedding net for discrete and continuous data.
+        embedding_net=combined_embedding_net,
+        num_bins=num_bins,
+        num_transforms=num_transforms,
+        num_components=num_components,
+        tail_bound=tail_bound,
+        hidden_features=hidden_features,
+        **kwargs,
+    )
+
+    return MixedDensityEstimator(
+        discrete_net=discrete_net,
+        continuous_net=continuous_net,
+        embedding_net=embedding_net_y,  # pass embedding for continuous condition part.
+        log_transform_input=log_transform_x,
+        input_shape=batch_x[0].shape,
+        condition_shape=batch_y[0].shape,
+    )
+
+
+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural likelihood estimator.
+
+    This estimator models p(x|theta) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (data), used to infer dimensionality.
+        batch_y: Batch of ys (parameters), used to infer dimensionality.
+        **kwargs: Additional arguments passed to _build_mixed_density_estimator.
+
+    Returns:
+        MixedDensityEstimator for MNLE.
+    """
+    return _build_mixed_density_estimator(
+        batch_x=batch_x, batch_y=batch_y, mode="mnle", **kwargs
+    )
+
+
+def build_mnpe(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural posterior estimator.
+
+    This estimator models p(theta|x) where x contains both continuous and discrete data.
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: dgedon
body: True, we could remove both functions and just use `_build_mixed_density_estimator`. 

The swapping of x/theta is happening because the function is once called as `likelihood_nn` and once as `posterior_nn`.
created_at: 2025-03-19T07:29:44Z
author_association: COLLABORATOR
diff_hunk: @@ -0,0 +1,151 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+from typing import Any, Callable, Dict, Optional, Union
+
+from torch.distributions import Distribution
+
+from sbi.inference.posteriors import (
+    DirectPosterior,
+    ImportanceSamplingPosterior,
+    MCMCPosterior,
+    RejectionPosterior,
+    VIPosterior,
+)
+from sbi.inference.trainers.npe.npe_c import NPE_C
+from sbi.neural_nets.estimators import MixedDensityEstimator
+from sbi.sbi_types import TensorboardSummaryWriter, TorchModule
+from sbi.utils.sbiutils import del_entries
+
+
+class MNPE(NPE_C):
+    def __init__(
+        self,
+        prior: Optional[Distribution] = None,
+        density_estimator: Union[str, Callable] = "mnpe",
+        device: str = "cpu",
+        logging_level: Union[int, str] = "WARNING",
+        summary_writer: Optional[TensorboardSummaryWriter] = None,
+        show_progress_bars: bool = True,
+    ):
+        r"""Mixed Neural Posterior Estimation (MNPE).
+
+        Like NPE-C, but designed to be applied to data with mixed types, e.g.,
+        continuous parameters and discrete parameters like they occur in models with
+        switching components. The emebedding net will only operate on the continuous
+        parameters, note this to design the dimension of the embedding net.
+
+        Args:
+            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()`.
+            density_estimator: If it is a string, it must be "mnpe" to use the
+                preconfigured neural nets for MNPE. Alternatively, a function
+                that builds a custom neural network can be provided. The function 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
+                needs to return a PyTorch `nn.Module` implementing the density
+                estimator. The density estimator needs to provide the methods
+                `.log_prob`, `.log_prob_iid()` and `.sample()`.
+            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.
+            summary_writer: A tensorboard `SummaryWriter` to control, among others, log
+                file location (default is `<current working directory>/logs`.)
+            show_progress_bars: Whether to show a progressbar during simulation and
+                sampling.
+        """
+
+        if isinstance(density_estimator, str):
+            assert (
+                density_estimator == "mnpe"
+            ), f"""MNPE can be used with preconfigured 'mnpe' density estimator only,
+                not with {density_estimator}."""
+        kwargs = del_entries(locals(), entries=("self", "__class__"))
+        super().__init__(**kwargs)
+
+    def train(
+        self,
+        training_batch_size: int = 200,
+        learning_rate: float = 5e-4,
+        validation_fraction: float = 0.1,
+        stop_after_epochs: int = 20,
+        max_num_epochs: int = 2**31 - 1,
+        clip_max_norm: Optional[float] = 5.0,
+        resume_training: bool = False,
+        discard_prior_samples: bool = False,
+        retrain_from_scratch: bool = False,
+        show_train_summary: bool = False,
+        dataloader_kwargs: Optional[Dict] = None,
+    ) -> MixedDensityEstimator:
+        density_estimator = super().train(
+            **del_entries(locals(), entries=("self", "__class__"))
+        )
+        assert isinstance(
+            density_estimator, MixedDensityEstimator
+        ), f"""Internal net must be of type
+            MixedDensityEstimator but is {type(density_estimator)}."""
+        return density_estimator
+
+    def build_posterior(
+        self,
+        density_estimator: Optional[TorchModule] = None,
+        prior: Optional[Distribution] = None,
+        sample_with: str = "direct",
+        mcmc_method: str = "slice_np_vectorized",
+        vi_method: str = "rKL",
+        direct_sampling_parameters: Optional[Dict[str, Any]] = None,
+        mcmc_parameters: Optional[Dict[str, Any]] = None,
+        vi_parameters: Optional[Dict[str, Any]] = None,
+        rejection_sampling_parameters: Optional[Dict[str, Any]] = None,
+        importance_sampling_parameters: Optional[Dict[str, Any]] = None,
+    ) -> Union[
+        MCMCPosterior,
+        RejectionPosterior,
+        VIPosterior,
+        DirectPosterior,
+        ImportanceSamplingPosterior,
+    ]:
+        """Build posterior from the neural density estimator.
+
+        Args:
+            density_estimator: The density estimator that the posterior is based on.
+                If `None`, use the latest neural density estimator that was trained.
+            prior: Prior distribution.
+            sample_with: Method to use for sampling from the posterior. Must be one of
+                [`direct` | `mcmc` | `rejection` | `vi` | `importance`].
+            mcmc_method: Method used for MCMC sampling, one of `slice_np`, `slice`,
+                `hmc`, `nuts`. Currently defaults to `slice_np` for a custom numpy
+                implementation of slice sampling; select `hmc`, `nuts` or `slice` for
+                Pyro-based sampling.
+            vi_method: Method used for VI, one of [`rKL`, `fKL`, `IW`, `alpha`].
+            direct_sampling_parameters: Additional kwargs passed to `DirectPosterior`.
+            mcmc_parameters: Additional kwargs passed to `MCMCPosterior`.
+            vi_parameters: Additional kwargs passed to `VIPosterior`.
+            rejection_sampling_parameters: Additional kwargs passed to `
+                RejectionPosterior`.
+            importance_sampling_parameters: Additional kwargs passed to
+                `ImportanceSamplingPosterior`.
+
+        Returns:
+            Posterior $p(\theta|x)$ with `.sample()` and `.log_prob()` methods.
+        """
+        if density_estimator is not None:
+            assert isinstance(
+                density_estimator, MixedDensityEstimator
+            ), f"""net must be of type MixedDensityEstimator but is {
+                type(density_estimator)
+            }."""
path: sbi/inference/trainers/npe/mnpe.py
user: login: dgedon
body: Good point. I'll add this additionally to the assertion
created_at: 2025-03-19T08:36:49Z
author_association: COLLABORATOR
path: sbi/neural_nets/net_builders/mnle.py
user: login: janfb
body: I don't think this was function was every explicitly user-facing. it was just used by `likelihood_nn` internally and the user would import `build_mnle' from `neural_nets.__init__.py`. So it might be fine to just remove this entirely here. 
created_at: 2025-03-19T08:54:46Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,18 +717,60 @@ def mcmc_transform(
                 constraint = prior.support.base_constraint  # type: ignore
             else:
                 constraint = prior.support
-            if isinstance(constraint, constraints._Real):
-                support_is_bounded = False
-            else:
-                support_is_bounded = True
-        else:
-            support_is_bounded = False
path: sbi/utils/sbiutils.py
user: login: janfb
body: why is this removed? 
created_at: 2025-03-19T08:55:31Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,18 +717,60 @@ def mcmc_transform(
                 constraint = prior.support.base_constraint  # type: ignore
             else:
                 constraint = prior.support
-            if isinstance(constraint, constraints._Real):
-                support_is_bounded = False
-            else:
-                support_is_bounded = True
-        else:
-            support_is_bounded = False
 
-        # Prior with bounded support, e.g., uniform priors.
-        if has_support and support_is_bounded:
-            transform = biject_to(prior.support)
-        # For all other cases build affine transform with mean and std.
+            # Check if the support is discrete
+            if hasattr(prior.support, "is_discrete"):
+                is_discrete = prior.support.is_discrete  # type: ignore
+            else:
+                is_discrete = False
+
+            if is_discrete:
+                # For discrete distributions, use mean/std transform
+                try:
+                    prior_mean = prior.mean.to(device)
+                    prior_std = prior.stddev.to(device)
+                except (NotImplementedError, AttributeError):
+                    warnings.warn(
+                        "The passed discrete prior has no mean or stddev attribute, "
+                        "estimating them from samples to build affine standardizing "
+                        "transform.",
+                        stacklevel=2,
+                    )
+                    theta = prior.sample(torch.Size((num_prior_samples_for_zscoring,)))
+                    prior_mean = theta.mean(dim=0).to(device)
+                    prior_std = theta.std(dim=0).to(device)
path: sbi/utils/sbiutils.py
user: login: janfb
body: move this into a small function to avoid code duplication? 
created_at: 2025-03-19T08:58:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -718,18 +717,60 @@ def mcmc_transform(
                 constraint = prior.support.base_constraint  # type: ignore
             else:
                 constraint = prior.support
-            if isinstance(constraint, constraints._Real):
-                support_is_bounded = False
-            else:
-                support_is_bounded = True
-        else:
-            support_is_bounded = False
path: sbi/utils/sbiutils.py
user: login: janfb
body: just moved down, ok. 
created_at: 2025-03-19T08:59:05Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,150 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    theta_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
path: tests/mnpe_test.py
user: login: janfb
body: this should be an `x_embedding` to avoid confusion
created_at: 2025-03-19T09:05:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,150 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    theta_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=theta_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
path: tests/mnpe_test.py
user: login: janfb
body: 👍 
created_at: 2025-03-19T09:05:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,150 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    theta_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=theta_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=5)
path: tests/mnpe_test.py
user: login: janfb
body: ```
max_num_epochs=1
```
to speed up tests.
created_at: 2025-03-19T09:05:52Z
author_association: CONTRIBUTOR
path: tutorials/Example_01_DecisionMakingModel.ipynb
user: login: janfb
body: please remove all diffs here. we probably want to have some kind of tutorial or how-to-guide with `MNPE`, but let's wait for the new documentation setup. 
created_at: 2025-03-19T09:07:16Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,150 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
path: tests/mnpe_test.py
user: login: janfb
body: can we come up with a simulator for which we have ground truth posterior to make a test on the accuracy? 

e.g., a mixture of Gaussians where the discrete params just selects the component: 

z ~ Categorical
x ~ N(mu_z, 1)

and Gaussian priors on mu_z with different prior means, e.g., -1 and 1. 

then for a fixed z, the posterior will just be a Gaussian composed from the likelihood and prior Gaussian just like in our standard linear Gaussian example.
created_at: 2025-03-19T09:15:07Z
author_association: CONTRIBUTOR
path: sbi/neural_nets/net_builders/mnle.py
user: login: dgedon
body: Possibly you wanted to mark not this but the following comment in sbiutil.py as "outdated"? 
Makes sense to remove `mnle.py`, I think
created_at: 2025-03-19T10:38:42Z
author_association: COLLABORATOR
diff_hunk: @@ -67,10 +67,14 @@ def build_mnle(
     hidden_features: int = 50,
     hidden_layers: int = 2,
     tail_bound: float = 10.0,
-    log_transform_x: bool = False,
+    log_transform_x: bool = True,
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: Sorry, I think I suggested setting this to `True`, but it works only for non-negative data / parameters should it should actually be set to `False` for `mnle` as well. 
created_at: 2025-03-20T09:53:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -123,41 +127,46 @@ def build_mnle(
         kwargs: additional keyword arguments passed to the flow model.
 
     Returns:
-        MixedDensityEstimator: nn.Module for performing MNLE.
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
     """
-
     check_data_device(batch_x, batch_y)
 
     warnings.warn(
-        "The mixed neural likelihood estimator assumes that x contains "
+        "The mixed neural density estimator assumes that x contains "
         "continuous data in the first n-k columns (e.g., reaction times) and "
         "categorical data in the last k columns (e.g., corresponding choices). If "
         "this is not the case for the passed `x` do not use this function.",
         stacklevel=2,
     )
+
     # Separate continuous and discrete data.
     if num_categorical_columns is None:
         num_disc = int(torch.sum(_is_discrete(batch_x)))
     else:
         num_disc = len(num_categorical_columns)
     cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
 
-    # Set up y-embedding net with z-scoring.
+    # The embeedding net is applied to the continuous part of the inputs
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
     if z_score_y_bool:
-        embedding_net = nn.Sequential(
+        embedding_net_y = nn.Sequential(
             standardizing_net(batch_y, structured_y), embedding_net
         )
-    # Get size of the embedded condition data.
-    embedded_batch_y = embedding_net(batch_y)
+    else:
+        embedding_net_y = embedding_net
+    embedding_net_x = nn.Identity()
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: we can remove this line, no? 
for both methods, `x` should not be embedded because it is learned.
created_at: 2025-03-20T09:56:40Z
author_association: CONTRIBUTOR
diff_hunk: @@ -123,41 +127,46 @@ def build_mnle(
         kwargs: additional keyword arguments passed to the flow model.
 
     Returns:
-        MixedDensityEstimator: nn.Module for performing MNLE.
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
     """
-
     check_data_device(batch_x, batch_y)
 
     warnings.warn(
-        "The mixed neural likelihood estimator assumes that x contains "
+        "The mixed neural density estimator assumes that x contains "
         "continuous data in the first n-k columns (e.g., reaction times) and "
         "categorical data in the last k columns (e.g., corresponding choices). If "
         "this is not the case for the passed `x` do not use this function.",
         stacklevel=2,
     )
+
     # Separate continuous and discrete data.
     if num_categorical_columns is None:
         num_disc = int(torch.sum(_is_discrete(batch_x)))
     else:
         num_disc = len(num_categorical_columns)
     cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
 
-    # Set up y-embedding net with z-scoring.
+    # The embeedding net is applied to the continuous part of the inputs
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
     if z_score_y_bool:
-        embedding_net = nn.Sequential(
+        embedding_net_y = nn.Sequential(
             standardizing_net(batch_y, structured_y), embedding_net
         )
-    # Get size of the embedded condition data.
-    embedded_batch_y = embedding_net(batch_y)
+    else:
+        embedding_net_y = embedding_net
+    embedding_net_x = nn.Identity()
+
+    # embed
+    embedded_batch_x = embedding_net_x(cont_x)
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: this should just be `batch_x` or `cont_x`
created_at: 2025-03-20T09:57:01Z
author_association: CONTRIBUTOR
diff_hunk: @@ -176,12 +185,15 @@ def build_mnle(
 
     # Set up a flow for modelling the continuous data, conditioned on the discrete data.
     continuous_net = model_builders[flow_model](
+        # TODO: add support for optional log-transform in flow builders.
         batch_x=(
-            torch.log(cont_x) if log_transform_x else cont_x
+            torch.log(embedded_batch_x + 1e-10)
+            if log_transform_x
+            else embedded_batch_x  # MNLE thing since x is data in mnle
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: change comment to sth like: use log-transform in case of strictly positive `x` (in principle this could apply to mnpe as well
created_at: 2025-03-20T09:58:09Z
author_association: CONTRIBUTOR
diff_hunk: @@ -195,8 +207,58 @@ def build_mnle(
     return MixedDensityEstimator(
         discrete_net=discrete_net,
         continuous_net=continuous_net,
-        embedding_net=embedding_net,  # pass embedding for continuous condition part.
+        embedding_net=embedding_net_y,  # pass embedding for continuous condition part.
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: rename this to just `embedding_net`? 
created_at: 2025-03-20T09:59:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -195,8 +207,58 @@ def build_mnle(
     return MixedDensityEstimator(
         discrete_net=discrete_net,
         continuous_net=continuous_net,
-        embedding_net=embedding_net,  # pass embedding for continuous condition part.
+        embedding_net=embedding_net_y,  # pass embedding for continuous condition part.
         log_transform_input=log_transform_x,
         input_shape=batch_x[0].shape,
         condition_shape=batch_y[0].shape,
     )
+
+
+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    log_transform_x: bool = True,
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: as said above, this should actually be set to `False`
created_at: 2025-03-20T09:59:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -195,8 +207,58 @@ def build_mnle(
     return MixedDensityEstimator(
         discrete_net=discrete_net,
         continuous_net=continuous_net,
-        embedding_net=embedding_net,  # pass embedding for continuous condition part.
+        embedding_net=embedding_net_y,  # pass embedding for continuous condition part.
         log_transform_input=log_transform_x,
         input_shape=batch_x[0].shape,
         condition_shape=batch_y[0].shape,
     )
+
+
+def build_mnle(
+    batch_x: Tensor,
+    batch_y: Tensor,
+    log_transform_x: bool = True,
+    **kwargs,
+) -> MixedDensityEstimator:
+    """Returns a mixed neural likelihood estimator.
+
+    This estimator models p(x|theta) where x contains both continuous and discrete data.
+
+    Args:
+        batch_x: Batch of xs (data), used to infer dimensionality.
+        batch_y: Batch of ys (parameters), used to infer dimensionality.
+        log_transform_x: whether to apply a log-transform to x. This is by default true
+            because x is data in MNLE.
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: update docstring accordingly. 
created_at: 2025-03-20T10:00:26Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
path: tests/mnpe_test.py
user: login: janfb
body: do we still need this one? could be just use the Bernoulli-Gaussian simulator as well? 
created_at: 2025-03-20T10:11:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
path: tests/mnpe_test.py
user: login: janfb
body: ```suggestion
    trainer.append_simulations(theta, x).train(max_num_epochs=1)
```
created_at: 2025-03-20T10:11:47Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
path: tests/mnpe_test.py
user: login: janfb
body: or use torch.bernoulli instead? 
created_at: 2025-03-20T10:12:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
path: tests/mnpe_test.py
user: login: janfb
body: ```suggestion
@pytest.mark.parametrize("flow_model", ("mdn", "nsf", "zuko_nsf"))
```
maf will learn a Gaussian in 1D, and it's fine if just cover mdn, nsf (nflows) and zuko. 
created_at: 2025-03-20T10:14:07Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
path: tests/mnpe_test.py
user: login: janfb
body: I think this should be changed to `z_score_x` because now `x` is in the condition (compared to `mnle`), no? 

I am actually not sure why we are testing this here. could also be removed I think
created_at: 2025-03-20T10:15:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
path: tests/mnpe_test.py
user: login: janfb
body: explicitly use embedding net? 
```
@pytest.mark.parametrize("embedding_net", (None, FCEmbedding(1, 1))
```
created_at: 2025-03-20T10:16:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test different samplers
+    posterior = trainer.build_posterior(prior=prior)
+    posterior.set_default_x(x_o)
+    if isinstance(posterior, VIPosterior):
+        posterior.train().sample((1,))
+    elif isinstance(posterior, RejectionPosterior):
+        posterior.sample((1,))
+    else:
+        posterior.sample((1,))
+
+
+class MixedSimulator(torch.nn.Module):
+    def __init__(self, sigma2):
+        super().__init__()
+        self.sigma2 = sigma2
+
+    def forward(self, theta):
+        # return x~N(mu_z, tau^2)
+        mu = theta[:, :2]
+        z = theta[:, -1]
+        # Select mu based on z
+        mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
+        # Sample from N(mu_z, sigma2)
+        x = Normal(loc=mu_z, scale=torch.sqrt(self.sigma2)).rsample()
+        return x.reshape(-1, 1)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("flow_model", ("nsf", "zuko_nsf"))
+def test_mnpe_accuracy(
+    flow_model: str,
+):
+    """Test MNPE accuracy."""
+    num_simulations = 1000
+    num_posterior_samples = 1000
+
+    z_score_theta = "none"
+    use_embed_net = False
+
+    # Prio
+    mu_0 = torch.tensor([-1.0])
+    mu_1 = torch.tensor([1.0])
+    tau2 = torch.tensor([0.1])
+    prior_mu0 = Normal(loc=mu_0, scale=torch.sqrt(tau2))
+    prior_mu1 = Normal(loc=mu_1, scale=torch.sqrt(tau2))
+
+    q = torch.tensor(0.8)
+    prior_z = Bernoulli(probs=q * torch.ones(1))
+
+    prior = MultipleIndependent(
+        [prior_mu0, prior_mu1, prior_z],
+    )
+
+    # Simulator
+    sigma2 = torch.tensor(0.1)
+    simulator = MixedSimulator(sigma2=sigma2)
+
+    # observation
+    theta_o = prior.sample((1,))
+    x_o = simulator(theta_o)
+
+    ############################################################
+    # MNPE
+    num_simulations = 4000
+    num_posterior_samples = 1_000
+
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
path: tests/mnpe_test.py
user: login: janfb
body: why use an embedding net here? 
created_at: 2025-03-20T10:18:57Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test different samplers
+    posterior = trainer.build_posterior(prior=prior)
+    posterior.set_default_x(x_o)
+    if isinstance(posterior, VIPosterior):
+        posterior.train().sample((1,))
+    elif isinstance(posterior, RejectionPosterior):
+        posterior.sample((1,))
+    else:
+        posterior.sample((1,))
+
+
+class MixedSimulator(torch.nn.Module):
+    def __init__(self, sigma2):
+        super().__init__()
+        self.sigma2 = sigma2
+
+    def forward(self, theta):
+        # return x~N(mu_z, tau^2)
+        mu = theta[:, :2]
+        z = theta[:, -1]
+        # Select mu based on z
+        mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
+        # Sample from N(mu_z, sigma2)
+        x = Normal(loc=mu_z, scale=torch.sqrt(self.sigma2)).rsample()
+        return x.reshape(-1, 1)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("flow_model", ("nsf", "zuko_nsf"))
+def test_mnpe_accuracy(
+    flow_model: str,
+):
+    """Test MNPE accuracy."""
+    num_simulations = 1000
+    num_posterior_samples = 1000
+
+    z_score_theta = "none"
+    use_embed_net = False
+
+    # Prio
+    mu_0 = torch.tensor([-1.0])
+    mu_1 = torch.tensor([1.0])
+    tau2 = torch.tensor([0.1])
+    prior_mu0 = Normal(loc=mu_0, scale=torch.sqrt(tau2))
+    prior_mu1 = Normal(loc=mu_1, scale=torch.sqrt(tau2))
+
+    q = torch.tensor(0.8)
+    prior_z = Bernoulli(probs=q * torch.ones(1))
+
+    prior = MultipleIndependent(
+        [prior_mu0, prior_mu1, prior_z],
+    )
+
+    # Simulator
+    sigma2 = torch.tensor(0.1)
+    simulator = MixedSimulator(sigma2=sigma2)
+
+    # observation
+    theta_o = prior.sample((1,))
+    x_o = simulator(theta_o)
+
+    ############################################################
+    # MNPE
+    num_simulations = 4000
+    num_posterior_samples = 1_000
+
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
path: tests/mnpe_test.py
user: login: janfb
body: can be removed I think
created_at: 2025-03-20T10:20:29Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test different samplers
+    posterior = trainer.build_posterior(prior=prior)
+    posterior.set_default_x(x_o)
+    if isinstance(posterior, VIPosterior):
+        posterior.train().sample((1,))
+    elif isinstance(posterior, RejectionPosterior):
+        posterior.sample((1,))
+    else:
+        posterior.sample((1,))
+
+
+class MixedSimulator(torch.nn.Module):
+    def __init__(self, sigma2):
+        super().__init__()
+        self.sigma2 = sigma2
+
+    def forward(self, theta):
+        # return x~N(mu_z, tau^2)
+        mu = theta[:, :2]
+        z = theta[:, -1]
+        # Select mu based on z
+        mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
+        # Sample from N(mu_z, sigma2)
+        x = Normal(loc=mu_z, scale=torch.sqrt(self.sigma2)).rsample()
+        return x.reshape(-1, 1)
+
+
+@pytest.mark.slow
+@pytest.mark.parametrize("flow_model", ("nsf", "zuko_nsf"))
+def test_mnpe_accuracy(
+    flow_model: str,
+):
+    """Test MNPE accuracy."""
+    num_simulations = 1000
+    num_posterior_samples = 1000
+
+    z_score_theta = "none"
+    use_embed_net = False
+
+    # Prio
+    mu_0 = torch.tensor([-1.0])
+    mu_1 = torch.tensor([1.0])
+    tau2 = torch.tensor([0.1])
+    prior_mu0 = Normal(loc=mu_0, scale=torch.sqrt(tau2))
+    prior_mu1 = Normal(loc=mu_1, scale=torch.sqrt(tau2))
+
+    q = torch.tensor(0.8)
+    prior_z = Bernoulli(probs=q * torch.ones(1))
+
+    prior = MultipleIndependent(
+        [prior_mu0, prior_mu1, prior_z],
+    )
+
+    # Simulator
+    sigma2 = torch.tensor(0.1)
+    simulator = MixedSimulator(sigma2=sigma2)
+
+    # observation
+    theta_o = prior.sample((1,))
+    x_o = simulator(theta_o)
+
+    ############################################################
+    # MNPE
+    num_simulations = 4000
+    num_posterior_samples = 1_000
+
+    theta = prior.sample((num_simulations,))
+    x = simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train()
+
+    # Test different samplers
+    posterior_mnpe = trainer.build_posterior(prior=prior)
+    posterior_mnpe.set_default_x(x_o)
+
+    posterior_mnpe_samples = posterior_mnpe.sample((num_posterior_samples,))
+
+    ############################################################
+    # True posterior
+    ###### Posterior of z ######
+    s2 = tau2 + sigma2
path: tests/mnpe_test.py
user: login: janfb
body: I think this should live in a separate function, we might want to re-use it somewhere else and the test would become more readable. 

created_at: 2025-03-20T10:21:52Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
+@pytest.mark.parametrize("use_embed_net", (True, False))
+def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
+    """Test MNPE API."""
+
+    # Generate data
+    num_simulations = 100
+    theta_true = torch.tensor([[0.5, 1.0]])
+    x_o = mixed_param_simulator(theta_true)
+
+    # Train and infer
+    prior = MultipleIndependent([
+        BoxUniform(torch.zeros(1), torch.ones(1)),  # continuous
+        Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+    ])
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    # Build estimator manually
+    x_embedding = FCEmbedding(1, 1)  # simple embedding net, 1 continuous parameter
+    density_estimator = posterior_nn(
+        model="mnpe",
+        flow_model=flow_model,
+        z_score_theta=z_score_theta,
+        embedding_net=x_embedding if use_embed_net else torch.nn.Identity(),
+        log_transform_x=False,
+    )
+    trainer = MNPE(density_estimator=density_estimator)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test different samplers
+    posterior = trainer.build_posterior(prior=prior)
+    posterior.set_default_x(x_o)
+    if isinstance(posterior, VIPosterior):
+        posterior.train().sample((1,))
+    elif isinstance(posterior, RejectionPosterior):
+        posterior.sample((1,))
+    else:
+        posterior.sample((1,))
+
+
+class MixedSimulator(torch.nn.Module):
+    def __init__(self, sigma2):
+        super().__init__()
+        self.sigma2 = sigma2
+
+    def forward(self, theta):
path: tests/mnpe_test.py
user: login: janfb
body: is there a reason to make this a nn.Module? 

couldn't it be just a `def mixed_param_simulator(...)`?

created_at: 2025-03-20T10:22:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -0,0 +1,282 @@
+# This file is part of sbi, a toolkit for simulation-based inference. sbi is licensed
+# under the Apache License Version 2.0, see <https://www.apache.org/licenses/>
+
+import pytest
+import torch
+from torch import Tensor
+from torch.distributions import Bernoulli, Normal
+
+from sbi.inference import MNPE
+from sbi.inference.posteriors.rejection_posterior import RejectionPosterior
+from sbi.inference.posteriors.vi_posterior import VIPosterior
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.embedding_nets import FCEmbedding
+from sbi.utils import BoxUniform
+from sbi.utils.torchutils import process_device
+from sbi.utils.user_input_checks_utils import MultipleIndependent
+from tests.test_utils import check_c2st
+
+
+# toy simulator for continuous data with mixed parameters
+def mixed_param_simulator(theta: Tensor) -> Tensor:
+    """Simulator for continuous data with mixed parameters.
+
+    Args:
+        theta: Parameters with mixed types - continuous and discrete.
+    Returns:
+        x: Continuous observation.
+    """
+    device = theta.device
+
+    # Extract parameters
+    a, b = theta[:, 0], theta[:, 1]
+    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
+    return (a + 2 * b).reshape(-1, 1) + noise
+
+
+@pytest.mark.gpu
+@pytest.mark.parametrize("device", ("cpu", "gpu"))
+def test_mnpe_on_device(
+    device,
+    num_simulations: int = 100,
+):
+    """Test MNPE API on device (CPU/GPU)."""
+    num_samples = 10
+    device = process_device(device)
+
+    # Generate data with mixed parameter types
+    theta_true = torch.tensor([[0.5, 1.0]], device=device)
+    x_o = mixed_param_simulator(theta_true).to(device)
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+        device=device,
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta).to(device)
+
+    trainer = MNPE(prior=prior, device=device)
+    trainer.append_simulations(theta, x).train(max_num_epochs=3)
+
+    # Test sampling on device
+    posterior = trainer.build_posterior()
+    samples = posterior.sample(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    assert samples.shape == (num_samples, 2)
+
+
+def test_batched_sampling(num_simulations: int = 100):
+    """Test MNPE API with batched sampling."""
+    batch_size = 5
+    num_samples = 10
+
+    # Generate data with mixed parameter types
+    theta_true = torch.cat(
+        (
+            torch.rand(batch_size, 1),
+            torch.ones(batch_size, 1),
+        ),
+        dim=1,
+    )
+    x_o = mixed_param_simulator(theta_true)  # This will return batch_size observations
+
+    prior = MultipleIndependent(
+        [
+            BoxUniform(low=torch.zeros(1), high=torch.ones(1)),  # continuous
+            Bernoulli(probs=0.8 * torch.ones(1)),  # discrete (Bernoulli)
+        ],
+    )
+    theta = prior.sample((num_simulations,))
+    x = mixed_param_simulator(theta)
+
+    trainer = MNPE(prior=prior)
+    trainer.append_simulations(theta, x).train(max_num_epochs=1)
+
+    # Test batched sampling
+    posterior = trainer.build_posterior()
+    samples = posterior.sample_batched(
+        sample_shape=(num_samples,),
+        x=x_o,
+        show_progress_bars=False,
+    )
+    print(samples.shape)
+    assert samples.shape == (num_samples, batch_size, 2)
+
+
+@pytest.mark.parametrize("flow_model", ("mdn", "maf", "nsf", "zuko_nsf"))
+@pytest.mark.parametrize("z_score_theta", ("independent", "none"))
path: tests/mnpe_test.py
user: login: dgedon
body: Since `z_score_theta` is tested in mnle, I think we should test `z_score_x` in here
created_at: 2025-03-20T11:01:16Z
author_association: COLLABORATOR
diff_hunk: @@ -79,21 +80,21 @@ def test_batched_sampling(num_simulations: int = 100):
     # Generate data with mixed parameter types
     theta_true = torch.cat(
         (
-            torch.rand(batch_size, 1),
-            torch.ones(batch_size, 1),
+            torch.rand(batch_size, 2),
+            torch.bernoulli(0.8 * torch.ones(batch_size, 1)),
         ),
         dim=1,
     )
path: tests/mnpe_test.py
user: login: janfb
body: you could just do `prior.sample((1,))` no?
created_at: 2025-03-20T11:14:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -79,21 +80,21 @@ def test_batched_sampling(num_simulations: int = 100):
     # Generate data with mixed parameter types
     theta_true = torch.cat(
         (
-            torch.rand(batch_size, 1),
-            torch.ones(batch_size, 1),
+            torch.rand(batch_size, 2),
+            torch.bernoulli(0.8 * torch.ones(batch_size, 1)),
         ),
         dim=1,
     )
path: tests/mnpe_test.py
user: login: janfb
body: we usually call it `theta_o` :) 
created_at: 2025-03-20T11:14:41Z
author_association: CONTRIBUTOR
diff_hunk: @@ -151,87 +152,23 @@ def test_mnpe_api(flow_model: str, z_score_theta: str, use_embed_net: bool):
         posterior.sample((1,))
 
 
-class MixedSimulator(torch.nn.Module):
-    def __init__(self, sigma2):
-        super().__init__()
-        self.sigma2 = sigma2
-
-    def forward(self, theta):
-        # return x~N(mu_z, tau^2)
-        mu = theta[:, :2]
-        z = theta[:, -1]
-        # Select mu based on z
-        mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
-        # Sample from N(mu_z, sigma2)
-        x = Normal(loc=mu_z, scale=torch.sqrt(self.sigma2)).rsample()
-        return x.reshape(-1, 1)
-
-
-@pytest.mark.slow
-@pytest.mark.parametrize("flow_model", ("nsf", "zuko_nsf"))
-def test_mnpe_accuracy(
-    flow_model: str,
-):
-    """Test MNPE accuracy."""
-    num_simulations = 1000
-    num_posterior_samples = 1000
-
-    z_score_theta = "none"
-    use_embed_net = False
-
-    # Prio
-    mu_0 = torch.tensor([-1.0])
-    mu_1 = torch.tensor([1.0])
-    tau2 = torch.tensor([0.1])
+def reference_posterior_mog(tau2, sigma2, mu_0, mu_1, x_o, q, num_posterior_samples):
+    """Reference posterior for a mixture of two Gaussians.
path: tests/mnpe_test.py
user: login: janfb
body: to avoid confusion with a classical (continuous MoG) I would call it a Bernoulli-Gaussian mixture
created_at: 2025-03-20T11:16:36Z
author_association: CONTRIBUTOR
diff_hunk: @@ -274,6 +211,53 @@ def test_mnpe_accuracy(
         (~z_mask).sum(),
     )).flatten()
 
+    return posterior_ref_samples
+
+
+@pytest.mark.slow
+def test_mnpe_accuracy():
path: tests/mnpe_test.py
user: login: janfb
body: great that this works so smoothly! 
created_at: 2025-03-20T11:22:42Z
author_association: CONTRIBUTOR
diff_hunk: @@ -18,20 +18,21 @@
 
 
 # toy simulator for continuous data with mixed parameters
-def mixed_param_simulator(theta: Tensor) -> Tensor:
+def mixed_simulator(theta, sigma2=0.1):
     """Simulator for continuous data with mixed parameters.
-
-    Args:
-        theta: Parameters with mixed types - continuous and discrete.
-    Returns:
-        x: Continuous observation.
-    """
+    Returns x~N(mu_z, tau^2)"""
     device = theta.device
 
-    # Extract parameters
-    a, b = theta[:, 0], theta[:, 1]
-    noise = 0.05 * torch.randn(a.shape, device=device).reshape(-1, 1)
-    return (a + 2 * b).reshape(-1, 1) + noise
+    if isinstance(sigma2, float):
+        sigma2 = torch.tensor(sigma2, device=device)
+
+    mu = theta[:, :2]
+    z = theta[:, -1]
+    # Select mu based on z
+    mu_z = torch.where(z == 0, mu[:, 0], mu[:, 1])
+    # Sample from N(mu_z, sigma2)
+    x = Normal(loc=mu_z, scale=torch.sqrt(sigma2)).rsample()
+    return x.reshape(-1, 1).to(device)
 
 
 @pytest.mark.gpu
path: tests/mnpe_test.py
user: login: janfb
body: this seems to be failing because of the MixedIndependent prior not allowing the move to GPU. Thus, I would suggest to just mark those tests as `xfail` (plus, we can actually remove the "cpu" test because this is tested below anyways):

`@pytest.mark.parametrize("device", (pytest.param("gpu", marks=pytest.mark.xfail(raises=NotImplementedError, reason="MultipleIndependent does not support device handling yet"))`


created_at: 2025-03-20T11:25:51Z
author_association: CONTRIBUTOR
diff_hunk: @@ -123,41 +127,44 @@ def build_mnle(
         kwargs: additional keyword arguments passed to the flow model.
 
     Returns:
-        MixedDensityEstimator: nn.Module for performing MNLE.
+        MixedDensityEstimator: nn.Module for performing MNLE or MNPE.
     """
-
     check_data_device(batch_x, batch_y)
 
     warnings.warn(
-        "The mixed neural likelihood estimator assumes that x contains "
+        "The mixed neural density estimator assumes that x contains "
         "continuous data in the first n-k columns (e.g., reaction times) and "
         "categorical data in the last k columns (e.g., corresponding choices). If "
         "this is not the case for the passed `x` do not use this function.",
         stacklevel=2,
     )
+
     # Separate continuous and discrete data.
     if num_categorical_columns is None:
         num_disc = int(torch.sum(_is_discrete(batch_x)))
     else:
         num_disc = len(num_categorical_columns)
     cont_x, disc_x = _separate_input(batch_x, num_discrete_columns=num_disc)
 
-    # Set up y-embedding net with z-scoring.
+    # The embeedding net is applied to the continuous part of the inputs
     z_score_y_bool, structured_y = z_score_parser(z_score_y)
     if z_score_y_bool:
         embedding_net = nn.Sequential(
             standardizing_net(batch_y, structured_y), embedding_net
         )
-    # Get size of the embedded condition data.
+    else:
+        embedding_net = embedding_net
+
+    # embed
     embedded_batch_y = embedding_net(batch_y)
     combined_condition = torch.cat([disc_x, embedded_batch_y], dim=-1)
 
     # Set up a categorical RV neural net for modelling the discrete data.
     discrete_net = build_categoricalmassestimator(
         disc_x,
         batch_y,
-        z_score_x="none",  # discrete data should not be z-scored.
-        z_score_y="none",  # y-embedding net already z-scores.
+        z_score_x="none",
+        z_score_y="none",
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: I suggest to keep those comments for explanation.
created_at: 2025-03-20T11:47:50Z
author_association: CONTRIBUTOR
diff_hunk: @@ -176,12 +183,15 @@ def build_mnle(
 
     # Set up a flow for modelling the continuous data, conditioned on the discrete data.
     continuous_net = model_builders[flow_model](
+        # TODO: add support for optional log-transform in flow builders.
         batch_x=(
-            torch.log(cont_x) if log_transform_x else cont_x
-        ),  # log transform manually.
+            torch.log(cont_x + 1e-10)
+            if log_transform_x
+            else cont_x  # can apply log transform if data is strictly positive
+        ),
         batch_y=combined_condition,
         z_score_x=z_score_x,
-        z_score_y="none",  # combined condition is already z-scored.
+        z_score_y="none",  # combined condition is already z-scored for mnle
path: sbi/neural_nets/net_builders/mixed_nets.py
user: login: janfb
body: applies to mnpe as well, no? just keep it as before? 
created_at: 2025-03-20T11:48:08Z
author_association: CONTRIBUTOR

PR comments:
user: login: dgedon
created_at: 2025-01-17T10:22:21Z
body: Update: 
- MNPE and tests are implemented
- for test with Bernoulli prior, I had to change mcmc_transforms to handle discrete distirbution. As default we just compute mean/std for discrete distributions
- currently MNPE with embedding nets does not work yet. Gives some backwards inplace operations error that I couldn't solve yet.
author_association: COLLABORATOR
user: login: janfb
created_at: 2025-02-25T18:59:53Z
body: @dgedon #1269 is now merged 🙌 
author_association: CONTRIBUTOR
user: login: dgedon
created_at: 2025-03-18T14:30:59Z
body: Updates:
- bug fix so everything works now. Essentially need to handle normalization with care when switching from mnle to mnpe
- remove unnecessary GPU handling (hackathon task). This limits MultipleIndependent as not allowing `device` argument yet
author_association: COLLABORATOR
user: login: dgedon
created_at: 2025-03-19T09:14:12Z
body: Update:
- simplify `_build_mixed_density_estimator` by not having `mode='mnpe'/'mnle'`
- add default `log_transform_x` as kwarg to `build_mnle` and `build_mnpe`
author_association: COLLABORATOR
user: login: dgedon
created_at: 2025-03-20T09:09:56Z
body: Update:
added an accuracy test for MNPE consisting of 2 Gaussians with varying mean. The observation is one Gaussian based on a Bernoulli "selection" variable. The analytic reference posterior is compared to the MNPE estimate using C2ST.
author_association: COLLABORATOR
user: login: codecov[bot]
created_at: 2025-03-20T10:10:01Z
body: ## [Codecov](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) Report
Attention: Patch coverage is `83.92857%` with `9 lines` in your changes missing coverage. Please review.
> Project coverage is 79.09%. Comparing base [(`46ccee0`)](https://app.codecov.io/gh/sbi-dev/sbi/commit/46ccee05eb065e50e5dc45dae4dbd4c16105c0d5?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) to head [(`bfaabe8`)](https://app.codecov.io/gh/sbi-dev/sbi/commit/bfaabe86492edb47f12790b503bf910479e8f338?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev).
> Report is 12 commits behind head on main.

| [Files with missing lines](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) | Patch % | Lines |
|---|---|---|
| [sbi/utils/sbiutils.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Futils%2Fsbiutils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL3V0aWxzL3NiaXV0aWxzLnB5) | 69.56% | [7 Missing :warning: ](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) |
| [sbi/inference/trainers/npe/mnpe.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Finference%2Ftrainers%2Fnpe%2Fmnpe.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL2luZmVyZW5jZS90cmFpbmVycy9ucGUvbW5wZS5weQ==) | 95.23% | [1 Missing :warning: ](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) |
| [sbi/utils/user\_input\_checks\_utils.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Futils%2Fuser_input_checks_utils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL3V0aWxzL3VzZXJfaW5wdXRfY2hlY2tzX3V0aWxzLnB5) | 50.00% | [1 Missing :warning: ](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) |

<details><summary>Additional details and impacted files</summary>


```diff
@@             Coverage Diff             @@
##             main    #1362       +/-   ##
===========================================
- Coverage   89.62%   79.09%   -10.53%     
===========================================
  Files         121      122        +1     
  Lines        9347     9394       +47     
===========================================
- Hits         8377     7430      -947     
- Misses        970     1964      +994     
```

| [Flag](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362/flags?src=pr&el=flags&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) | Coverage Δ | |
|---|---|---|
| [unittests](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362/flags?src=pr&el=flag&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) | `79.09% <83.92%> (-10.81%)` | :arrow_down: |

Flags with carried forward coverage won't be shown. [Click here](https://docs.codecov.io/docs/carryforward-flags?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#carryforward-flags-in-the-pull-request-comment) to find out more.

| [Files with missing lines](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev) | Coverage Δ | |
|---|---|---|
| [sbi/inference/\_\_init\_\_.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Finference%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL2luZmVyZW5jZS9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [sbi/inference/trainers/npe/\_\_init\_\_.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Finference%2Ftrainers%2Fnpe%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL2luZmVyZW5jZS90cmFpbmVycy9ucGUvX19pbml0X18ucHk=) | `100.00% <100.00%> (ø)` | |
| [.../neural\_nets/estimators/mixed\_density\_estimator.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Fneural_nets%2Festimators%2Fmixed_density_estimator.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL25ldXJhbF9uZXRzL2VzdGltYXRvcnMvbWl4ZWRfZGVuc2l0eV9lc3RpbWF0b3IucHk=) | `94.73% <ø> (-1.76%)` | :arrow_down: |
| [sbi/neural\_nets/factory.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Fneural_nets%2Ffactory.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL25ldXJhbF9uZXRzL2ZhY3RvcnkucHk=) | `90.90% <100.00%> (ø)` | |
| [sbi/neural\_nets/net\_builders/\_\_init\_\_.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Fneural_nets%2Fnet_builders%2F__init__.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL25ldXJhbF9uZXRzL25ldF9idWlsZGVycy9fX2luaXRfXy5weQ==) | `100.00% <100.00%> (ø)` | |
| [sbi/neural\_nets/net\_builders/mixed\_nets.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Fneural_nets%2Fnet_builders%2Fmixed_nets.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL25ldXJhbF9uZXRzL25ldF9idWlsZGVycy9taXhlZF9uZXRzLnB5) | `97.05% <100.00%> (ø)` | |
| [sbi/inference/trainers/npe/mnpe.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Finference%2Ftrainers%2Fnpe%2Fmnpe.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL2luZmVyZW5jZS90cmFpbmVycy9ucGUvbW5wZS5weQ==) | `95.23% <95.23%> (ø)` | |
| [sbi/utils/user\_input\_checks\_utils.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Futils%2Fuser_input_checks_utils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL3V0aWxzL3VzZXJfaW5wdXRfY2hlY2tzX3V0aWxzLnB5) | `88.67% <50.00%> (-1.13%)` | :arrow_down: |
| [sbi/utils/sbiutils.py](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362?src=pr&el=tree&filepath=sbi%2Futils%2Fsbiutils.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev#diff-c2JpL3V0aWxzL3NiaXV0aWxzLnB5) | `78.38% <69.56%> (-9.16%)` | :arrow_down: |

... and [38 files with indirect coverage changes](https://app.codecov.io/gh/sbi-dev/sbi/pull/1362/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=sbi-dev)

</details>
<details><summary>🚀 New features to boost your workflow: </summary>

- ❄ [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
author_association: NONE