New file structure for neural networks (#1237)

* New file structure for neural networks

* no more imports of the builders from sbi.neural_nets

* Some fixups for paths

* Rename build_functions to build_nets

* rename to net_builders

* raise ImportError for embedding nets

Author: michaeldeistler

sbi/inference/fmpe/fmpe_base.py

@@ -15,7 +15,8 @@
 from sbi import utils as utils
 from sbi.inference.base import NeuralInference
 from sbi.inference.posteriors.direct_posterior import DirectPosterior
-from sbi.neural_nets import ConditionalDensityEstimator, flowmatching_nn
+from sbi.neural_nets import flowmatching_nn
+from sbi.neural_nets.estimators import ConditionalDensityEstimator
 from sbi.utils import (
     RestrictedPrior,
     handle_invalid_x,

sbi/inference/potentials/likelihood_based_potential.py

@@ -8,12 +8,14 @@
 from torch.distributions import Distribution
 
 from sbi.inference.potentials.base_potential import BasePotential
-from sbi.neural_nets.estimators import ConditionalDensityEstimator
+from sbi.neural_nets.estimators import (
+    ConditionalDensityEstimator,
+    MixedDensityEstimator,
+)
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
     reshape_to_sample_batch_event,
 )
-from sbi.neural_nets.mnle import MixedDensityEstimator
 from sbi.sbi_types import TorchTransform
 from sbi.utils.sbiutils import mcmc_transform
 

sbi/inference/snle/mnle.py

@@ -9,7 +9,7 @@
 from sbi.inference.posteriors import MCMCPosterior, RejectionPosterior, VIPosterior
 from sbi.inference.potentials import mixed_likelihood_estimator_based_potential
 from sbi.inference.snle.snle_base import LikelihoodEstimator
-from sbi.neural_nets.mnle import MixedDensityEstimator
+from sbi.neural_nets.estimators import MixedDensityEstimator
 from sbi.sbi_types import TensorboardSummaryWriter, TorchModule
 from sbi.utils.sbiutils import del_entries
 from sbi.utils.user_input_checks import check_prior

sbi/inference/snle/snle_base.py

@@ -16,7 +16,8 @@
 from sbi.inference.posteriors import MCMCPosterior, RejectionPosterior, VIPosterior
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.potentials import likelihood_estimator_based_potential
-from sbi.neural_nets import ConditionalDensityEstimator, likelihood_nn
+from sbi.neural_nets import likelihood_nn
+from sbi.neural_nets.estimators import ConditionalDensityEstimator
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
 )

sbi/inference/snpe/snpe_base.py

@@ -24,7 +24,8 @@
 from sbi.inference.posteriors.base_posterior import NeuralPosterior
 from sbi.inference.posteriors.importance_posterior import ImportanceSamplingPosterior
 from sbi.inference.potentials import posterior_estimator_based_potential
-from sbi.neural_nets import ConditionalDensityEstimator, posterior_nn
+from sbi.neural_nets import posterior_nn
+from sbi.neural_nets.estimators import ConditionalDensityEstimator
 from sbi.neural_nets.estimators.shape_handling import (
     reshape_to_batch_event,
     reshape_to_sample_batch_event,

sbi/neural_nets/__init__.py

@@ -1,26 +1,27 @@
-from sbi.neural_nets.classifier import (
-    build_linear_classifier,
-    build_mlp_classifier,
-    build_resnet_classifier,
-)
-from sbi.neural_nets.embedding_nets import (
-    CNNEmbedding,
-    FCEmbedding,
-    PermutationInvariantEmbedding,
-)
-from sbi.neural_nets.estimators import ConditionalDensityEstimator, NFlowsFlow
 from sbi.neural_nets.factory import (
     classifier_nn,
     flowmatching_nn,
     likelihood_nn,
     posterior_nn,
+    posterior_score_nn,
 )
-from sbi.neural_nets.flow import (
-    build_made,
-    build_maf,
-    build_maf_rqs,
-    build_nsf,
-    build_zuko_maf,
-)
-from sbi.neural_nets.mdn import build_mdn
-from sbi.neural_nets.mnle import MixedDensityEstimator, build_mnle
+
+
+def __getattr__(name):
+    if name in ["CNNEmbedding", "FCEmbedding", "PermutationInvariantEmbedding"]:
+        raise ImportError(
+            "As of sbi v0.23.0, you have to import embedding networks from "
+            "`sbi.neural_nets.embedding_nets`. For example, use: "
+            f"`from sbi.neural_nets.embedding_nets import {name}`"
+        )
+    elif name == "classifier_nn":
+        return classifier_nn
+    elif name == "flowmatching_nn":
+        return flowmatching_nn
+    elif name == "likelihood_nn":
+        return likelihood_nn
+    elif name == "posterior_nn":
+        return posterior_nn
+    elif name == "posterior_score_nn":
+        return posterior_score_nn
+    raise AttributeError(f"Module '{__name__}' has no attribute '{name}'")