DEV Community: Satwik Sai Prakash Sahoo

Wrapping Up My GSoC 2026 Journey with sbi

Satwik Sai Prakash Sahoo — Mon, 24 Aug 2026 02:31:36 +0000

This is my final work product for Google Summer of Code 2026, where I spent the summer redesigning how sbi builds neural networks.

What the project is about?

sbi does Bayesian inference for simulators you cannot write a likelihood for. You have a simulator and a prior, you run simulations, and a neural network learns the posterior over the simulator's parameters from those runs. It implements a whole family of methods (NPE, NLE, NRE, FMPE, NPSE and the mixed variants), and every one of them trains some kind of neural density estimator.

My project was about how you pick that network and configure it. Over the summer I replaced the old string and factory-function interface with typed configuration objects, one class per model, across every estimator family in the library. The goal was that a setting the chosen model cannot use fails immediately when you write it, rather than being quietly dropped on the way to the network, and that existing code keeps working through a deprecation path instead of breaking.

Organisation: NumFOCUS | Sub-organisation: sbi-dev/sbi | Mentors: Jan Teusen (@janfb) & Nicholas Junge (@nicholasjng)

What was the problem?

The old API had three ways to specify a network, and they had grown separately:

A string: NPE(prior, density_estimator="nsf"). Easy, but no way to set hyperparameters.
A factory function: posterior_nn(model="nsf", hidden_features=64). It returned an opaque closure. You could not inspect it, print it or serialize it, and it accepted any keyword you gave it.
A custom module, for full control.

The factory functions were the problem. They spanned a whole family of models, so their signature was the union of every setting any model might want. A setting that the chosen model did not use was quietly discarded. There was no type information anywhere, so editors could not autocomplete and type checkers could not help.

The goal was a layered API where each layer is typed and inspectable, mistakes fail at configuration time instead of at training time, and existing code keeps working through a deprecation path rather than breaking.

How is the problem solved now?

Every estimator family now has typed config objects:

from sbi.inference import NPE, NRE
from sbi.neural_nets import ZukoNSFConfig, ResNetClassifierConfig

trainer = NPE(prior, density_estimator=ZukoNSFConfig(hidden_features=64, num_transforms=8))
trainer = NRE(prior, classifier=ResNetClassifierConfig(hidden_features=64, num_blocks=3))

They are frozen dataclasses, so they print cleanly, you can log them in an experiment record, and your editor autocompletes the fields. Strings still work and emit a FutureWarning that names the class to switch to. Custom modules still work.

The part I am most happy with is what happens when you are wrong. There are three distinct failures, all before training:

NSFConfig(hiden_features=64)          # TypeError  - misspelled name
NSFConfig(z_score_input="strucured")  # ValueError - misspelled value
MAFConfig(num_bins=8)                 # TypeError  - not a setting this model has

What each PR did

#1872 laid the groundwork. The shared base class, the supporting types, and a protocol rename to free up the name "Builder" for the new objects.

#1877 added the first real builder and #1882 wired it into the NPE trainers, with the deprecation path for strings.

The most useful thing that happened here was something we removed. The original design had build() take a BuildContext object carrying shape, device and z-scoring statistics. While implementing it, Jan noticed the parameter was never actually read: everything it held could be derived from the training batch that was already being passed in. We cut it. It was a premature abstraction, and recognising that early saved a lot of machinery that would have needed maintaining.

#1904 extended the builder to NLE, and this is where the design got interesting.

NPE models the parameters given the data. NLE models the data given the parameters. Same builder, opposite roles. A signature of build(batch_theta, batch_x) reads correctly for one and backwards for the other, which is exactly how you get a bug where the standardization lands on the wrong variable.

So build() became build(batch_input, batch_condition). input is whatever the model is modelling and condition is whatever it is conditioning on, and the trainer decides which is which. The config never has to know. That decision carried cleanly through every family that came after, including the vector-field one that landed months later.

#1912 added the mixed density estimators for MNLE and MNPE, where part of the data is discrete.

#1920 added the NRE classifiers and bundled in a set of fixes to the shared base class. The reasoning was timing: this was the first PR proving the base served a third family, so fixing it there meant NRE and everything after inherited the fixes instead of needing a retrofit.

What went in:

Invalid Literal values now raise at construction. A misspelled field name already raised. A misspelled value on a correctly named field did not, and stayed silent until training.
Model-incompatible settings raise, by inspecting the target build function's signature.
frozen=True on the configs, so they are immutable.
A custom __repr__ that prints only what you actually set.
Uniform public API exposure, since one builder had been left unexported.

It also carried a rename I am glad we did. The configs used z_score_x and z_score_y, inherited from the low-level functions. Those names are ambiguous across families: for NPE, z_score_x standardizes the parameters, not x. They became z_score_input and z_score_condition, matching the role-neutral signature. The builder API had not shipped in a release yet, so the rename was free. A few weeks later it would have been a breaking change.

#1921 covered FMPE and NPSE. This was also the point where the working style changed. Instead of implementing a design handed to me, I started writing a short design proposal first and getting it reviewed before writing code. It felt slower for about a day and then obviously correct: a design mistake caught in a document costs an hour, and the same mistake caught in code review costs a week.

Partway through, we changed direction, and I think it made the whole API better. Until then, each family had one flat class holding every field any model in that family might want. To stop you setting a field your model does not use, I had built machinery: applicability lists, signature inspection, and drift tests to keep those lists honest as models changed.

The new idea deletes all of it by making the class itself the choice. One small config per model, carrying only the fields that model accepts. Then there is nothing to validate, because a field that does not apply simply does not exist, and Python's own argument checking raises for free.

#1975 piloted that on the marginal trainer with nine classes, one per Zuko flow. #1986 then converted the density, classifier and mixed families, replacing three flat builders with 15 density configs, 3 classifier configs and one mixed config.

Where things stand?

Merged and in sbi:

PR	What
#1872	Foundation types and renames
#1877	`DensityEstimatorBuilder` with `build()` dispatch
#1882	Integration into the NPE trainers
#1904	Integration into NLE, role-neutral signature
#1912	`MixedDensityEstimatorBuilder` for MNLE and MNPE
#1920	`RatioEstimatorBuilder` for NRE, base class hardening, z-score rename
#1921	`VectorFieldEstimatorBuilder` for FMPE and NPSE
#1975	Per-model configs for `MarginalTrainer`
#1986	Per-model configs for the density, classifier and mixed families

Open for review:

#1997 finishes the mixed family. The internal mixed build function was still taking two APIs at once, the old flat arguments and a config, so it now takes only a config and the deprecated strings build one in the factory. It also fixes a compatibility bug I would never have caught by counting parameters: the flat API passed a wider spline tail_bound to every model that reads it, and keying that off the model name instead of the field silently narrowed the tail for one model while leaving the parameter count identical.

Draft, queued behind each other:

#1996 converts the vector-field family to the per-model design. This one is the most interesting of the set. The builder picks along what looks like three axes at once (flow or score matching, which SDE, which network architecture), but the axes are not independent: the estimator type and the SDE together select exactly one class. So it becomes two composed objects, FlowMatchingConfig(net=MLPConfig(...)), which is seven classes instead of the sixteen you get from multiplying the axes out.
#1987 is the documentation PR: a new how-to for the config API, a rework of the "abstraction levels" guide, the config classes added to the API reference, and the deprecated strings cleared out of the tutorials. It stays a draft until #1996 merges, because it documents the API that PR lands, then it gets rebased and reviewed.

The stacking matters here. Each of these sits on the one before it, so they merge in order and each gets rebased onto main once its predecessor lands. That is a habit I picked up the hard way, and it is why the diffs stay readable.

What is left to do?

Test consolidation. Each PR added its own test file, which was right during development and is now redundant. We agreed to consolidate them into fewer parametrized files once everything merges. The distinction that matters for CI cost is that tests which train are the expensive ones, while construction and default checks are cheap, so the consolidation should keep the cheap ones broad and the expensive ones few.

Retiring the legacy factory functions. posterior_nn, likelihood_nn and classifier_nn still exist and still work. They now build the new configs internally, so they are already thin. Once the deprecation window closes they can go, along with the legacy *Config validators behind them.

The stretch goal: data-source-agnostic builds. This is the one I did not get to, and it is the most interesting remaining piece. Right now z-scoring statistics are computed inside the build functions from a single in-memory batch. For datasets too large to hold in one tensor, those statistics would have to come from a pre-pass over a dataloader, or be accumulated online. The plan is to move statistics and shapes out of the build functions and into a context object that the trainer fills in, which is the BuildContext we deliberately deferred back in phase 1. The private types for it are already in the codebase as a reserved seam, so it is a local addition rather than a re-architecture. It is correctness-sensitive, since wrong z-scoring silently degrades inference rather than erroring, and it touches roughly 24 build functions, which is why it stayed a stretch goal.

Acknowledgement

I am really grateful to my mentor Jan Teusen for giving me this opportunity and for reviews that consistently found the thing I had not thought about, and for pushing me to write designs before code in the second half. To Nicholas Junge for the backup mentorship. And to the sbi maintainers for treating a student's PRs with the same seriousness as anyone else's. This was a great learning experience for me.

The coding period is over but the work is not. I am continuing with sbi through November to land the vector-field conversion and the documentation, and I intend to stick around after that as a permanent contributor.

I have also written some blogs every two weeks throughout the program describing my learnings in detail.

Signing off,
Satwik Sai Prakash Sahoo
GitHub | LinkedIn | Bluesky | X

Final weeks of GSoC

Satwik Sai Prakash Sahoo — Wed, 19 Aug 2026 12:50:31 +0000

Hey everyone! Week 11 is done and week 12 is about to be over, which means the official GSoC coding period is ending soon. These two weeks brought a design change that I think makes the whole API better, a merged PR, and a couple of things about these flows that I only noticed because of the new design.

PR #1921 merged

First, the vector field builder from my last post is in. FMPE and NPSE now go through the same typed interface as everything else, so four of the five families are done.

A change of direction: one class per model

Up to this point every family had one flat builder class holding every field any model in that family might want. DensityEstimatorBuilder has fields for MAF, for NSF, for MDN, for MADE and for nine Zuko flows, all in one class. To stop people from setting a field their chosen model does not use, I had built machinery: applicability checks, whitelists, signature inspection, and drift tests to keep those lists honest.

The new idea is to delete all of that by making the class itself the choice. One small config class per model, carrying only the fields that model actually accepts:

MarginalTrainer(density_estimator=MarginalNSFConfig(num_transforms=8))

If a setting does not belong to that model, it is not a field on the class, so Python raises a TypeError at construction. No whitelist, no applicability check, no drift test to keep in sync. Python's own argument checking does the job for free.

This is was a major change in design decision, and PR #1975 pilots it on the marginal trainer.

PR #1975: per-model marginal configs

So this PR adds nine classes, one for each unconditional Zuko flow: BPF, GF, MAF, NAF, NCSF, NICE, NSF, SOSPF and UNAF. The trainer default also moves from a string to None, which resolves to MarginalNSFConfig(), so that a plain MarginalTrainer() does not trip the new deprecation warning. Strings still work and still warn.

I also parametrized the tests over all nine models instead of testing one and hoping the rest behaved, which caught things I would have missed.

Writing nine classes made me look at nine models

The thing I did not expect from this design is how much it forces you to actually read the model you are writing a class for. With one flat builder you write the union of every field and move on. With one class per model you have to answer, nine times, what does this flow accept and what does it not.

Two of them turned out to be genuinely different from the rest.

Gaussianization flows do not take a width. MarginalGFConfig has no hidden_features field, which looks like an oversight until you follow where the value goes. Zuko routes the settings it does not name itself into an element-wise transform, and that transform only builds a network when there is a condition to build it from. A marginal flow is unconditional, so there is never a condition, so the network is never built and the width is dropped on the way there. Rather than accept a number and quietly ignore it, the class refuses it and points you at num_transforms and components, which are the settings that actually change the flow.

Circular splines do not live on the real line. This one I found the hard way, through a CI failure. NCSF is a neural circular spline flow, so its base distribution is uniform on
[-pi, pi] rather than a standard normal. My test was drawing torch.randn samples for every model in the same loop, which is fine for eight of them and outside the domain for the ninth. The fix was a small helper that draws each model's data from the space it is actually defined on:

def _samples_in_domain_of(model):
    if model == "ncsf":
        return torch.pi * (2 * torch.rand(100, 3) - 1)
    return torch.randn(100, 3)

Neither of these is a bug in sbi. They are just facts about the models that the old flat class had no place to record, so nobody had to think about them. A per-model class has exactly the right place: the field list, and the docstring next to it.

That is the part of this design I have come to like most. It is not only that wrong settings fail faster. It is that writing the classes made the differences between the models visible, to me while writing them and to the next person reading them.

How I am using AI on this project

Something worth writing about, since it came up in both syncs.

I used Claude Code for the first time for the marginal PR, and disclosed it in the PR description, which is now our standard practice. The workflow we settled on is that a second model reviews the first one's output adversarially. Claude writes, then Codex reviews it critically, and it is surprisingly good at finding bugs and pointing at code that is more verbose than it needs to be.

But still the problem, I just started using Claude Code and although the code generation is fast, I sometimes have a hard time surfing around changed files to keep or remove AI generated code, the continuous changes and decisions somehow still take the same time when I wasn't using AI tools. But over the time I will get hold of it and finally it would increase my productivity. Jan also provided me some cool resources to learn Claude Code which I am very grateful for. Also the thing Jan said that stuck with me is that the early months of writing everything by hand, though slow, were the right foundation. The project understanding those months built is what makes the AI assisted work reviewable now. I can tell when a suggestion is wrong, and I have caught a few that were, including one that would have quietly swapped one flow implementation for a different one.

Tests and what happens now

I raised a concern about CI cost, since we now have a test file per PR and they keep growing. The answer was a useful distinction: tests that train are the expensive ones, while tests that just check defaults and construction are cheap. So the plan is to keep plenty of tests during development and decide before merge which ones to keep.

The official coding period ends this week, and the documentation PR is my final contribution for the evaluation. But the work does not stop there. I am carrying on with sbi through November to implement our stretch goals which would make the user experience even better.

What's Next?

The density and classifier configs are in a draft PR, the documentation PR comes after that, and then the vector field conversion. Four families down, one to go.

Thanks for sticking with me through the whole coding period, and see you in the next update!

Learnt Git Rebasing the hard way

Satwik Sai Prakash Sahoo — Wed, 19 Aug 2026 12:49:35 +0000

Hello again! Weeks 9 and 10 covered the fourth estimator family, the first release that carries my work, and a git mess that taught me more than the code did.

PR #1920 merged, and sbi 0.27 shipped

First, the ratio estimator builder from my last post made it in, together with the whole base class hardening bundle. Right after that, sbi 0.27 was released and PRs 1 to 5 were ported over to main. So the typed builder API is now actually in a release, for NPE, NLE, MNPE, MNLE and all four NRE variants.

That also means the free rename window is closed. The z_score_input and z_score_condition names we picked in week 7 are now the shipped names, which is exactly why we did that rename when we did.

Writing the design before the code

Under the new workflow I described last time, the vector field work started with a markdown file instead of a Python file.

The vector field family is genuinely harder than the others. build_vector_field_estimator picks along what looks like three axes at the same time: whether you want flow matching or score matching, which SDE type you want if it is score matching, and which network architecture sits inside. And unlike every other family, there are no per-model build functions to hang a config class on.

I wrote up the options with some open questions then my mentor, Jan reviewed it and we settled on one builder with a base class abstraction to cut the redundancy across builders. Doing this on paper first was clearly the right call. Some of the things I had assumed while writing the proposal turned out to be wrong, and finding that out in a review comment on a markdown file was a lot cheaper than finding it out in a review comment on 800 lines of code.

PR #1921: `VectorFieldEstimatorBuilder`

With the design agreed, the implementation covers FMPE and NPSE. The builder takes the architecture as model, one of mlp, ada_mlp, transformer or transformer_cross_attn, plus estimator_type for flow versus score and sde_type for the noise schedule.

One detail I liked: this family has fields that only make sense for one side. sigma_min, sigma_max and the schedule settings belong to score estimators, and gaussian_baseline belongs to flow matching. So the builder raises if you set a score field with estimator_type="flow", and the other way round.

Interestingly the raw build_vector_field_estimator function does not do that. It silently ignores sde_type if you asked for flow matching. So the builder is stricter than the function it wraps, which is the whole point of this project in one sentence.

I also kept the applicability lists deliberately small and added tests that check them against the existing config classes, so the two cannot drift apart without CI noticing.

The rebasing lesson

Now for the part that actually cost me time.

PR #1921 was branched off PR #1920, which was branched off earlier work. As PR #1920 changed during review, I kept merging main back into my branch to stay current. What I ended up with was a branch history full of merge commits and small repair commits, and when it came time to present a clean PR, replaying that history with git rebase was painful. Every conflict I had already resolved once came back to be resolved again.

The thing that finally worked, and what Jan uses too, is to stop trying to replay history at all:

git reset --soft <clean base>
git commit

--soft keeps your working tree exactly as it is and just moves the branch pointer, so you get the final state of your code sitting on a clean base, and then you commit it in whatever shape makes sense for a reviewer. All the intermediate mess disappears because you never replay it.

The other half of the lesson is that sbi squash merges its PRs. I had assumed that once PR #1920 merged, its commits would stop showing up in PR #1921's diff on their own. They do not. Squash merging creates a brand new commit on main, so my branch's merge base never moves and those commits keep appearing until I explicitly rebase onto the new main:

git rebase --onto upstream/main <old base> <my branch>

I now understand a lot better why maintainers care so much about branch hygiene. A reviewer opening a PR with 40 commits, half of which belong to a different PR, is going to have a bad time no matter how good the actual code is.

Also decided

Two smaller things from the week 9 sync. Keyword harmonization is moving to the documentation PR rather than the vector field PR, so the VF PR stays focused and we do not break tutorials in the same change. And the work order for the rest of the project is now: vector field PR, then test refactoring, then documentation.

What's Next?

PR #1921 is in review. After that comes the marginal estimator builder, and then the documentation PR. There is also a growing pile of test files, one per PR at this point, that we plan to consolidate once everything has merged.

Thanks for following along, and see you in the next one!

Passing GSoC Midterms

Satwik Sai Prakash Sahoo — Wed, 19 Aug 2026 12:41:19 +0000

Hey everyone! Weeks 7 and 8 are done, and the biggest news first: I passed the midterm evaluation🎉. Half of GSoC is behind me now.

These two weeks were less about writing new classes and more about going back and making the ones I already had a lot stricter. Here is what happened.

PR #1920: The `RatioEstimatorBuilder`

This PR adds the builder for the NRE family, so NRE_A, NRE_B, NRE_C and BNRE all get the same typed interface that NPE and NLE already had.

The builder itself was honestly the easy part. By now the pattern is well established, so it was mostly mirroring what already worked, with linear, mlp and resnet as the classifier options. The interesting part was what my mentor Jan Teusen suggested we bundle into the same PR.

Some improvements to the base class

PR #1920 was the first PR that proved the shared base class serves a third family. My Mentor, Jan pointed out that this was exactly the right moment to fix the validation gaps in that base, because then the NRE builder and the vector field builder that was coming next would inherit the fixes for free, instead of me retrofitting four builders later.

So we folded a hardening bundle into the same PR:

Invalid Literal values now raise at construction. This was the real gap. If you typed a wrong field name, Python already raised a TypeError for you. But if you typed a wrong value on a correctly named field, like z_score_input="idependent", nothing happened until you called .train() and it blew up much later. Now it fails immediately.

Model-incompatible kwargs now raise too. Passing num_blocks to a linear classifier used to be silently dropped, and the same argument on mlp crashed late. Both are caught at construction now by inspecting the target build function's signature.

frozen=True on the config dataclasses. Configs are now immutable. If you want a different setting, you make a new object instead of mutating the old one. This sounds like a small thing but it removes a whole class of "I changed the config after passing it to the trainer, why did nothing happen" confusion.

A custom __repr__. Printing a config now shows only the fields that matter for the chosen model, instead of a wall of thirty Nones.

Uniform public API exposure. The mixed density builder was not exported at the same level as the others, so I fixed that too.

The z-score rename

The other thing that landed in this PR was a rename I had been putting off. The builders used z_score_x and z_score_y, inherited from the low level build_* functions. The problem is that those names are ambiguous once you have more than one family. For NPE, z_score_x actually standardizes theta, not x, because NPE models theta given x. So the name says one thing and does another depending on which trainer you are in.

We renamed them to z_score_input and z_score_condition. input is whatever the model is modelling and condition is whatever it is conditioning on, and the trainer decides which is which. The downstream build_* functions keep their old names, and the builder translates at the boundary.

The reason we did it now and not later is simple: the builder API has never shipped in a release yet, so this rename is free. After a release it becomes a breaking change. It was in three builders at that point and PR #1921 was about to add a fourth.

A shift in how we work

The other outcome of the midterm sync was a change in workflow for the second half of the project. Instead of Jan handing me a design and me implementing it, I now write a short design proposal first, as a markdown file in our GSoC repo, and Jan reviews it before I write any code.

The first one of these is for the vector field estimators, which is the next family. I will be honest, writing a design document before touching the editor feels slower at first. But the point is that a design mistake caught in a markdown file costs an hour, and the same mistake caught in review costs a week of rework. We also decided to start the documentation earlier than originally planned, instead of leaving it all to the very end.

An honest note about week 8

Week 8 was a slow week for me. I had other commitments and reduced my workload, so PR #1920 did not move much and the vector field design proposal did not get written. I flagged it in the sync rather than quietly falling behind, and Jan was completely fine about it. Back to normal pace from now.

What's Next?

PR #1920 needs the hardening review addressed and then it can merge. After that I write the vector field design proposal, get it reviewed, and only then start on VectorFieldEstimatorBuilder for FMPE and NPSE.

Thanks for reading, and see you in the next update!

Halfway Through GSoC

Satwik Sai Prakash Sahoo — Mon, 20 Jul 2026 05:33:56 +0000

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 $p (θ ∣ x)$ , meaning $θ$ is the input and $x$ is the condition.

However, NLE estimates the likelihood $p (x ∣ θ)$ . In this case, $x$ 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!

First internal integration of the new API

Satwik Sai Prakash Sahoo — Mon, 29 Jun 2026 17:47:42 +0000

Hello again! The GSoC coding period is in full swing, and weeks 3 and 4 have been absolutely packed with progress.

Following up on the foundational work from my first PR, I have just successfully merged two major PRs (#1877 and #1882) into the gsoc-2026 branch for sbi. These PRs introduce the first concrete builder class and wire it directly into the core Neural Posterior Estimation (NPE) trainers.

Here is a dive into what I built, the technical hurdles, and a very valuable lesson I learned about software architecture along the way.

PR #1877: The `DensityEstimatorBuilder`

The primary goal of this phase was to replace the old, opaque posterior_nn() and likelihood_nn() factory closures with something typed, inspectable, and much more robust.

To solve this, I introduced the DensityEstimatorBuilder. It inherits from the base contract we established in Week 1 and serves as the unified entry point for creating neural networks in sbi. Using the __post_init__ method in Python dataclasses, it immediately validates the model name against a _VALID_DENSITY_MODELS set, failing early if the user provides an unknown architecture.

Major architectural shift

Initially, our plan dictated that the build() method should take a BuildContext object. The idea was that this context would hold all necessary information, including pre-computed z-scoring stats, and pass it neatly down the chain.

However, as I implemented the body of the build() method, my mentor Jan Teusen noticed that the context parameter wasn't actually being used. Every piece of information the builder needed could be derived directly from the raw batch_theta and batch_x tensors.

We realized that forcing the BuildContext into this signature was a premature abstraction.

Instead of holding onto a design just because it was the original plan, we decided to defer the context object entirely until the z-scoring stats are actually pre-computed in a later phase. We updated our roadmap and simplified the build signature.

def build(self, batch_theta: Tensor, batch_x: Tensor):
    pass

It was a great decision to defer the context implementation to work on the breadth first and then move towards depth.

PR #1882: Integrating the Builder into NPE Trainers

With the builder merged, the next step was integration. I updated PosteriorEstimatorTrainer, NPE_B, NPE_C, and MNPE to accept the new DensityEstimatorBuilder instead of relying solely on strings or callables.

To maintain backward compatibility while moving the API forward, a graceful deprecation path was implemented. If a user passes a string (e.g., "maf"), the code still works perfectly, but it now emits a FutureWarning.

if density_estimator is None:
    self._build_neural_net = self._wrap_builder(DensityEstimatorBuilder(model="maf"))
elif isinstance(density_estimator, str):
    warnings.warn(
        "Passing a string for `density_estimator` is deprecated. "
        "Use DensityEstimatorBuilder(model=...) instead.",
        FutureWarning,
        stacklevel=3,
    )
    self._build_neural_net = posterior_nn(model=density_estimator)
elif isinstance(density_estimator, _EstimatorBuilderBase):
    self._build_neural_net = self._wrap_builder(density_estimator)
else:
    self._build_neural_net = density_estimator

Overal feedback on my work

The code review for PR #1882 was intense but incredibly rewarding. My mentor provided feedback on how to write tests that are not just concise, but strong and explicit in their intent.

For example, I originally wrote a test that checked if passing a callable avoided triggering the deprecation warning. But I wasn't actually asserting that no warning was thrown, I was just running the code and assuming silence meant success.

My mentor showed me how to use warnings.catch_warnings() with a strict filter to instantly fail the test if a FutureWarning leaked through:

import warnings

with warnings.catch_warnings():
    warnings.simplefilter("error", FutureWarning)
    inference = NPE_C(prior, density_estimator=builder, show_progress_bars=False)

We also did a deep dive into correct type hinting and managing default arguments. I initially left density_estimator="maf" as the default argument in the NPE_C initialization. My mentor pointed out that this would cause the deprecation warning to fire every single time a user initialized the class without arguments! The fix was to change the type hint default to None and handle the "maf" fallback inside the logic block.

What's Next?

Weeks 3 and 4 were a massive leap forward for the API refactor. We now have a working, integrated builder that correctly handles all continuous density estimators.

Next up, I will be tackling the remaining likelihood and classifier builders. Thanks for following along on this journey, and see you in the next update!

Merging My First PR for GSoC!

Satwik Sai Prakash Sahoo — Mon, 15 Jun 2026 13:12:48 +0000

Hey everyone! The official coding period for Google Summer of Code (GSoC) has begun, and I am thrilled to share that my very first Pull Request (PR) for the sbi (Simulation-Based Inference) repository has been officially merged!

These first two weeks have been packed with writing code, writing tests, and most importantly, going through an incredibly insightful code review process with my mentors.

Here is a breakdown of what I worked on, the technical challenges I faced, and the best practices I learned along the way.

Laying the Foundation (PR #1872)

My GSoC project is focused on refactoring the Neural Network (NN) Builder API. Before we can build the new, shiny network builders, we need a solid foundation. This first PR was all about setting up the necessary data structures and renaming existing protocols to make room for the new architecture.

1. The New `build_context.py`

I created a new file to house the core pieces needed to set up a neural network. This centralizes how the data is prepared before it hits the network:

ZScoreConfig: Tracks how the user wants to preprocess data.
ZScoreStats: Holds the calculated mean and standard deviation for the data.
BuildContext: Bundles everything required to build a network (shapes, device, dtype, and z-score stats) into a single, clean dataclass.
compute_z_score_stats(): A helper function to calculate these statistics directly from the training tensors.

2. Clearing the Naming Space

To make the codebase more intuitive, I had to rename a few core components across 16 different trainer files:

Changed ConditionalEstimatorBuilder to ConditionalEstimatorBuildFn. This clarifies that the protocol is actually a function, not an object, and frees up the "Builder" name for upcoming classes.
Changed _EstimatorConfigBase to _EstimatorBuilderBase, giving it an empty build() method to prepare for the next phase.

Big Takeaways from Code Review

Getting the code working was only half the battle. The review process with my mentor, taught me several advanced Python and PyTorch practices.

Managing Tensor Devices

When creating the BuildContext, the data (x and theta) are passed in as PyTorch tensors. My mentor pointed out that if x and theta accidentally end up on different devices (e.g., one on CPU, one on GPU), it will cause runtime failures. I updated the code to explicitly check that both tensors share the same device and added GPU-specific pytest runs to ensure everything handles device placement gracefully.

Dataclasses and Tensor Equality

I originally set up my statistics classes as @dataclass(frozen=True). However, because they hold PyTorch tensors, this creates a subtle bug. PyTorch's implementation of equality (==) for tensors returns a boolean tensor, not a single boolean value, which breaks the dataclass equality checks.

The fix was simple but crucial:

@dataclass(frozen=True, eq=False)
class ZScoreStats:
    theta_mean: Optional[Tensor] = None
    theta_std: Optional[Tensor] = None
    x_mean: Optional[Tensor] = None
    x_std: Optional[Tensor] = None

By adding eq=False, equality becomes identity-based, completely bypassing the PyTorch tensor comparison crash.

Graceful Deprecation with PEP 562

Since ConditionalEstimatorBuilder was used by the community, we couldn't just delete the name and break everyone's code. Instead of leaving a dummy class, I learned how to use module-level __getattr__ to intercept the import and throw a FutureWarning.

def __getattr__(name: str):
    if name == "ConditionalEstimatorBuilder":
        import warnings
        warnings.warn(
            "`ConditionalEstimatorBuilder` has been renamed to "
            "`ConditionalEstimatorBuildFn`. The old name still works but will be "
            "removed in a future release. Update your import to: "
            "`from sbi.neural_nets.estimators.base import ConditionalEstimatorBuildFn`.",
            FutureWarning,
            stacklevel=2,
        )
        return ConditionalEstimatorBuildFn
    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")

Refactoring Tests

I used Claude Opus to help generate initial test coverage for the new context types. It gave me class-based tests. However, the sbi standard relies heavily on function-based tests and pytest fixtures. I completely refactored the test suite to use @pytest.mark.parametrize to test valid combinations of inputs, ensuring maximum coverage with minimal, highly readable code.

What's Next?

With the foundation types successfully merged into the gsoc-2026 branch, the groundwork is officially laid out! Up next is PR #1877, where I will be adding the actual DensityEstimatorBuilder with dynamic build() dispatching.

Stay tuned for week 3 and 4 updates!

GSoC Community Bonding Period: Getting Ready to Code

Satwik Sai Prakash Sahoo — Thu, 04 Jun 2026 06:51:17 +0000

Hey everyone! Welcome back to my Google Summer of Code (GSoC) journey.

In my last post, I shared the story of how I got into open source and was selected for GSoC with NumFOCUS to work on the Neural Network Builder API Refactor project for sbi (Simulation-Based Inference).

Since the official announcement, the past three weeks have been dedicated to the Community Bonding Period. It is designed to help contributors get to know their mentors, understand the community culture, and familiarize themselves with the codebase and tools.

Here is exactly what I did during these past three weeks to get ready for the main coding phase!

The Kickoff Meeting

We started the bonding period with a great kickoff call on Google Meet. It was a joint meeting that included the mentors for both of the selected sbi projects, the selected GSoC candidates. We were also joined by the mentee who successfully completed the GSoC project for sbi last year!

Everyone introduced themselves, and it was incredibly inspiring to meet the team face-to-face (virtually!) and hear about everyone's backgrounds. Having a former GSoC student there was a huge bonus, as they shared some great insights into what to expect in the coming months.

Setting Up the Machine

A big part of getting started is making sure the development environment is properly configured. During our meetings, we discussed the machine setup in detail to ensure both candidates had everything required to run and test the sbi codebase locally without any hiccups.

Embracing AI Coding Assistants

One of the most interesting discussions we had was about using AI coding assistants. In the modern development world, tools like these are becoming standard, and our mentors actually encouraged us to use them!

However, they emphasized using them carefully and strictly following project guidelines. To help us get the most out of these tools without compromising code quality, the mentors shared some excellent Claude code tutorials and provided us with resources to level up our AI coding skills. It is all about using AI to augment our capabilities, not to blindly write code for us.

Diving into the Codebase and Planning

With the setup out of the way, the rest of the three weeks was all about reading. I spent a lot of time diving deep into the sbi codebase and reading through the documentation properly.

Understanding the architecture is critical for my API refactor project. To keep everything organized and on track, we were provided with a plan.md file. This document serves as our roadmap, outlining exactly what needs to be tackled, the skills required, and the milestones we need to hit.

What's Next?

The community bonding period was incredibly helpful. It gave me the time to get comfortable with the team, set up my workflow, and build a solid foundation before the heavy lifting begins.

Now, the official coding phase is here! I am beyond excited to start pushing commits and working on the Neural Network Builder API Refactor.

Thanks for reading, and stay tuned. I will be sharing more technical updates and lessons learned as I start writing code for GSoC!

How I Got Selected for GSoC: My Journey with sbi

Satwik Sai Prakash Sahoo — Sun, 31 May 2026 19:55:27 +0000

Hey everyone! I am super excited to share that I have been selected for Google Summer of Code (GSoC) with NumFOCUS. I will be working on the Neural Network Builder API Refactor project for the organization sbi (Simulation-Based Inference).

If you are wondering how to get started with open source or what the GSoC process is really like, here is the story of how I got here. Stay tuned, because I will keep updating you on my GSoC journey and sharing my experiences as I get deeper into the project!

Why I Started Open Source

Being a fan of Linus Torvalds, I was always fascinated by open source. Also, I wanted to test my coding skills. I wanted to know: Is my code actually good enough for the real world? In my first year of college, I spent a lot of time doing competitive programming. At the same time, I was teaching myself machine learning and mathematics, working mostly with Python and PyTorch.

Finding the Right Project

When I decided to try open source, I wanted to find an organization that actually matched what I was learning. Because I was already into machine learning and PyTorch, I went looking for something in that space.

That is how I found sbi.

Making my first Pull Request (PR) was a little scary, but the maintainers at sbi were incredibly welcoming. They patiently guided me through my early, beginner-level PRs. Thanks to their feedback, I learned so much about how to write clean, high-quality code.

The GSoC Surprise

After a few months of fixing bugs and contributing, I found out from the maintainers that sbi would be participating in GSoC this year. They encouraged me and appreciated the work I had done over the past few months, for which I will be forever grateful.

But there was one problem: I hadn't researched GSoC at all. I had no idea what the requirements were or how to write a project proposal.

Getting Help with the Proposal

Writing the proposal was tough since it was my first time. Luckily, I didn't have to do it alone.

The maintainers also stepped up to mentor me. They reviewed my proposal drafts, gave me detailed feedback, and helped me improve it step by step until it was ready.

What's Next?

All that hard work and community support paid off, and I got accepted!

The GSoC community bonding period just finished up. In my next blog post, I will share exactly what I did during this bonding period to get ready for the main coding phase.

Thanks for reading, and stay tuned!

DEV Community: Satwik Sai Prakash Sahoo

Wrapping Up My GSoC 2026 Journey with sbi

What the project is about?

What was the problem?

How is the problem solved now?

What each PR did

Where things stand?

What is left to do?

Acknowledgement

Final weeks of GSoC

PR #1921 merged

A change of direction: one class per model

PR #1975: per-model marginal configs

Writing nine classes made me look at nine models

How I am using AI on this project

Tests and what happens now

What's Next?

Learnt Git Rebasing the hard way

PR #1920 merged, and sbi 0.27 shipped

Writing the design before the code

PR #1921: VectorFieldEstimatorBuilder

The rebasing lesson

Also decided

What's Next?

Passing GSoC Midterms

PR #1920: The RatioEstimatorBuilder

Some improvements to the base class

The z-score rename

A shift in how we work

An honest note about week 8

What's Next?

Halfway Through GSoC

PR #1904: Integrating the Builder into NLE

The Role-Neutral build() Signature

Parameterizing Tests

PR #1912: Mixed Density Estimators (MNLE & MNPE)

"continuous_model" over "model"

Failing Fast is Better Than Failing Late

Catching Configuration Drift

What's Next?

First internal integration of the new API

PR #1877: The DensityEstimatorBuilder

Major architectural shift

PR #1882: Integrating the Builder into NPE Trainers

Overal feedback on my work

What's Next?

Merging My First PR for GSoC!

Laying the Foundation (PR #1872)

1. The New build_context.py

2. Clearing the Naming Space

Big Takeaways from Code Review

Managing Tensor Devices

Dataclasses and Tensor Equality

Graceful Deprecation with PEP 562

Refactoring Tests

What's Next?

GSoC Community Bonding Period: Getting Ready to Code

The Kickoff Meeting

Setting Up the Machine

Embracing AI Coding Assistants

Diving into the Codebase and Planning

What's Next?

How I Got Selected for GSoC: My Journey with sbi

Why I Started Open Source

Finding the Right Project

The GSoC Surprise

Getting Help with the Proposal

What's Next?

PR #1921: `VectorFieldEstimatorBuilder`

PR #1920: The `RatioEstimatorBuilder`

The Role-Neutral `build()` Signature

PR #1877: The `DensityEstimatorBuilder`

1. The New `build_context.py`