Simformer

docs/advanced_tutorials.rst

@@ -33,6 +33,7 @@ Neural nets
    advanced_tutorials/12_iid_data_and_permutation_invariant_embeddings.ipynb
    advanced_tutorials/19_flowmatching_and_scorematching.ipynb
    advanced_tutorials/20_score_based_methods_new_features.ipynb
+   advanced_tutorials/22_simformer.ipynb
 
 
 Training

docs/sbi.rst

@@ -29,6 +29,8 @@ Neural nets
    sbi.neural_nets.classifier_nn
    sbi.neural_nets.flowmatching_nn
    sbi.neural_nets.posterior_score_nn
+   sbi.neural_nets.simformer_score_nn
+   sbi.neural_nets.simformer_flow_nn
    sbi.neural_nets.marginal_nn
 
 
@@ -60,6 +62,8 @@ Training
    sbi.inference.MNPE
    sbi.inference.FMPE
    sbi.inference.NPSE
+   sbi.inference.Simformer
+   sbi.inference.FlowMatchingSimformer
    sbi.inference.NLE_A
    sbi.inference.NRE_A
    sbi.inference.NRE_B

sbi/diagnostics/misspecification.py

@@ -8,8 +8,7 @@
 from typing import Optional
 
 import torch
-import torch.nn as nn
-from torch import Tensor
+from torch import Tensor, nn
 
 from sbi.inference.trainers.npe.npe_base import PosteriorEstimatorTrainer
 from sbi.neural_nets.estimators import UnconditionalDensityEstimator
@@ -157,10 +156,9 @@ def calc_misspecification_mmd(
             )
         if inference._neural_net.embedding_net is None:
             raise AttributeError(
-                "embedding_net attribute is None but is required for misspecification"
-                " detection."
+                "embedding_net attribute is None but is required for misspecification "
+                "detection."
             )
-
         z_obs = inference._neural_net.embedding_net(x_obs).detach()
         z = inference._neural_net.embedding_net(x).detach()
     else:

sbi/inference/__init__.py

@@ -3,6 +3,7 @@
 
 from sbi.inference.abc import MCABC, SMCABC
 from sbi.inference.trainers.base import (
+    MaskedNeuralInference,  # noqa: F401
     NeuralInference,  # noqa: F401
     check_if_proposal_has_default_x,
     infer,
@@ -11,7 +12,7 @@
 from sbi.inference.trainers.nle import MNLE, NLE_A
 from sbi.inference.trainers.npe import MNPE, NPE_A, NPE_B, NPE_C  # noqa: F401
 from sbi.inference.trainers.nre import BNRE, NRE_A, NRE_B, NRE_C  # noqa: F401
-from sbi.inference.trainers.vfpe import FMPE, NPSE
+from sbi.inference.trainers.vfpe import FMPE, NPSE, FlowMatchingSimformer, Simformer
 
 SNL = SNLE = SNLE_A = NLE = NLE_A
 _nle_family = ["NLE"]
@@ -33,7 +34,13 @@
 _abc_family = ["ABC", "MCABC", "SMC", "SMCABC"]
 
 
-__all__ = _npe_family + _nre_family + _nle_family + _abc_family + ["FMPE", "NPSE"]
+__all__ = (
+    _npe_family
+    + _nre_family
+    + _nle_family
+    + _abc_family
+    + ["FMPE", "NPSE", "Simformer", "FlowMatchingSimformer"]
+)
 
 from sbi.inference.posteriors import (
     DirectPosterior,
@@ -53,4 +60,14 @@
 )
 from sbi.utils.simulation_utils import simulate_for_sbi
 
-__all__ = ["FMPE", "MarginalTrainer", "NLE", "NPE", "NPSE", "NRE", "simulate_for_sbi"]
+__all__ = [
+    "FlowMatchingSimformer",
+    "FMPE",
+    "MarginalTrainer",
+    "NLE",
+    "NPE",
+    "NPSE",
+    "NRE",
+    "Simformer",
+    "simulate_for_sbi",
+]

sbi/inference/posteriors/vector_field_posterior.py

@@ -25,7 +25,7 @@
 from sbi.samplers.score.predictors import Predictor
 from sbi.sbi_types import Shape
 from sbi.utils import check_prior
-from sbi.utils.sbiutils import gradient_ascent, within_support
+from sbi.utils.sbiutils import gradient_ascent, handle_invalid_x, within_support
 from sbi.utils.torchutils import ensure_theta_batched
 
 
@@ -202,6 +202,17 @@ def sample(
 
         x = self._x_else_default_x(x)
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
+        _, num_nans, num_infs = handle_invalid_x(x)
+
+        if num_nans + num_infs == 0:
+            ValueError(
+                "Some invalid entries (NaN/Infs) were "
+                "found in x. You probably passed these as the ground observed "
+                "x's `x_obs`. "
+                "Please, remove these values and provide reasonable observed x's "
+                "to avoid the sampling process to run indefinitely."
+            )
+
         is_iid = x.shape[0] > 1
         self.potential_fn.set_x(
             x,
@@ -455,6 +466,17 @@ def sample_batched(
         """
         num_samples = torch.Size(sample_shape).numel()
         x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
+        _, num_nans, num_infs = handle_invalid_x(x)
+
+        if num_nans + num_infs == 0:
+            ValueError(
+                "Some invalid entries (NaN/Infs) were "
+                "found in x. You probably passed these as the ground observed "
+                "x's `x_obs`. "
+                "Please, remove these values and provide reasonable observed x's "
+                "to avoid the sampling process to run indefinitely."
+            )
+
         condition_dim = len(self.vector_field_estimator.condition_shape)
         batch_shape = x.shape[:-condition_dim]
         batch_size = batch_shape.numel()