Hey everyone! We have officially hit the halfway mark of the Google Summer of Code (GSoC) 2026 coding period.
These past two weeks have been heavily focused on expanding the Neural Network Builder API beyond just standard NPE (Neural Posterior Estimation) models. I have successfully merged PRs #1904 and #1912 into the gsoc-2026 branch for sbi.
Here is a breakdown of what I worked on, the architectural challenges we solved, and some vital lessons learned about failing fast and role-neutral design.
PR #1904: Integrating the Builder into NLE
In my previous PR, we successfully wired the DensityEstimatorBuilder into NPE trainers. The next logical step was to do the same for Neural Likelihood Estimation (NLE) trainers. However, this brought up an immediate architectural challenge regarding how the builder signs its inputs.
The Role-Neutral build() Signature
Previously, the build() method took (batch_theta, batch_x). This made sense for NPE, which estimates the posterior distribution
, meaning
is the input and
is the condition.
However, NLE estimates the likelihood . In this case, is the input and is the condition.
If we kept the (batch_theta, batch_x) signature on the base builder, it would be extremely confusing and semantically incorrect for NLE.
The fix? Role-Neutral Signatures.
My mentor Jan Teusen suggested that I should refactor the base builder and DensityEstimatorBuilder methods to accept (batch_input, batch_condition).
Now, the trainers handle the role assignment via the _wrap_builder() closure:
-
For NPE:
build(batch_input=batch_theta, batch_condition=batch_x) -
For NLE:
build(batch_input=batch_x, batch_condition=batch_theta)
This change cleanly decouples the mathematical model the network represents from the data types the network consumes.
Parameterizing Tests
During review, my mentor also suggested that the NPE and NLE integration tests were essentially "twin" tests doing the exact same thing with different objects, we parameterized them!
TRAINERS = [(NPE_C, posterior_nn, "theta"), (NLE_A, likelihood_nn, "x")]
@pytest.mark.parametrize("trainer,factory", [(t, f) for t, f, _ in TRAINERS], ids=["npe", "nle"])
def test_string_emits_deprecation_warning(trainer, factory):
prior = MultivariateNormal(zeros(2), eye(2))
with pytest.warns(FutureWarning, match="deprecated"):
trainer(prior, density_estimator="maf", show_progress_bars=False)
This saved a ton of code duplication while keeping the coverage identical.
PR #1912: Mixed Density Estimators (MNLE & MNPE)
For datasets containing mixed data types (both continuous and discrete variables), sbi uses Mixed Neural Posterior/Likelihood Estimators (MNPE/MNLE).
To support these, I created a sibling class to our main builder: the MixedDensityEstimatorBuilder.
"continuous_model" over "model"
Our main design choice here was using the field continuous_model instead of just model. This is because the supported mixed architectures include non-flow models like MDN and MADE. Now, a user can explicitly configure the continuous sub-network: MixedDensityEstimatorBuilder(continuous_model="nsf").
Failing Fast is Better Than Failing Late
MNPE requires a MixedDensityEstimator. If a user accidentally passed a standard DensityEstimatorBuilder (e.g., MNPE(density_estimator=DensityEstimatorBuilder())), the code would gladly accept it, train an entire continuous normal flow for hundreds of epochs, and then crash at the very end when it finally checked the returned class type.
Failing at the end of a long training run with a cryptic error is terrible UX.
We added a strict type-check inside the trainer initialization to fail fast:
elif isinstance(density_estimator, _EstimatorBuilderBase) and not isinstance(
density_estimator, MixedDensityEstimatorBuilder
):
raise TypeError(
"MNPE requires a MixedDensityEstimatorBuilder; got "
f"{type(density_estimator).__name__}. Use "
"MixedDensityEstimatorBuilder(continuous_model=...)."
)
Now, if a user makes this mistake, the code crashes instantly with clear instructions on how to fix it, saving hours of compute time.
Catching Configuration Drift
We defined _VALID_MIXED_CONTINUOUS_MODELS to validate inputs. However, hardcoding this set means it could drift away from the actual supported models in the factory if another developer adds a model later. To prevent this, I added a test specifically designed to catch configuration drift:
def test_valid_continuous_models_match_builders():
from sbi.neural_nets.net_builders.mixed_nets import model_builders
assert _VALID_MIXED_CONTINUOUS_MODELS == frozenset(model_builders)
What's Next?
We are halfway there! The continuous and mixed density estimators are successfully wired into the system. Next up, we will be tackling the remaining likelihood and classifier builders to complete the NN Builder ecosystem.
Thanks for reading, and stay tuned for the next update!
Top comments (0)