Final weeks of GSoC

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!