Add RatioEstimatorBuilder for NRE trainers and base class improvements

[satwiksps]

What does this PR do?

This PR continues the GSoC 2026 Neural Network Builder API refactor by introducing the RatioEstimatorBuilder for all four NRE trainers (NRE_A, NRE_B, NRE_C, BNRE).

Instead of passing a string like classifier="resnet" and hiding settings in kwargs, we can now pass a fully typed and clear configuration like RatioEstimatorBuilder(model="resnet", hidden_features=64) to use linear, MLP, or ResNet models. Passing None will silently default to a ResNet builder to keep the old behavior, while using the old string arguments is now deprecated.

Files Changed

estimator_configs.py

• New RatioEstimatorBuilder dataclass configures and validates linear/MLP/ResNet classifiers, with build() routing to the right function.
• Builders are now frozen/immutable (use dataclasses.replace() to tweak), with safer extra_kwargs handling and a cleaner __repr__ that only shows what's actually set.
• Renamed z_score_x/y to z_score_input/condition for clarity, with aliasing so it still maps correctly downstream.
• Validates at construction time, catches bad Literal values and flags hyperparameters that don't apply to the chosen model.

nre_base.py

• classifier now takes a RatioEstimatorBuilder (defaults to None -> ResNet); old string configs still work but warn, wrong types fail fast.
• Added _wrap_builder to bridge the new builder API to the legacy (batch_theta, batch_x) format.

nre_a.py, nre_b.py, nre_c.py, bnre.py

• Same classifier type hint/default update, docstrings refreshed to match.

npe_base.py, nle_base.py, mnle.py, mnpe.py

• Deprecation warnings now point to the exact import path for the new builder, plus a small isinstance-check cleanup in mnle.py/mnpe.py.

sbi/neural_nets/__init__.py

• Exports DensityEstimatorBuilder and MixedDensityEstimatorBuilder so those import paths actually resolve.

density_estimator_builder_test.py

• Updated for the z_score_input/z_score_condition rename.

nre_builder_integration_test.py (new)

• Integration tests across all NRE trainers: defaults, legacy warnings, validation, immutability, repr, z-score aliasing.

npe_nle_builder_integration_test.py

• Keeps builder fields in sync with build-function signatures, plus z-score alias tests.

AI Usage

• Used Grammarly for PR description refinement, VS Code Copilot for code suggestions, and GPT OSS via Antigravity for docstrings and cosmetic updates.

Does this close any issues?

N/A

Any relevant code examples, logs, or error messages?

N/A

[codecov]

Codecov Report

❌ Patch coverage is 99.25926% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 88.01%. Comparing base (587f922) to head (8122159).
⚠️ Report is 31 commits behind head on gsoc-2026.

Files with missing lines	Patch %	Lines
sbi/neural_nets/net_builders/estimator_configs.py	99.02%	1 Missing ⚠️

Additional details and impacted files

@@              Coverage Diff              @@
##           gsoc-2026    #1920      +/-   ##
=============================================
+ Coverage      87.99%   88.01%   +0.01%     
=============================================
  Files            144      144              
  Lines          13499    13704     +205     
=============================================
+ Hits           11879    12061     +182     
- Misses          1620     1643      +23     

Flag	Coverage Δ
fast	81.67% <99.25%> (?)

Flags with carried forward coverage won't be shown. Click here to find out more.

Files with missing lines	Coverage Δ
sbi/inference/trainers/nle/mnle.py	97.14% <100.00%> (+4.55%)	⬆️
sbi/inference/trainers/nle/nle_base.py	97.36% <ø> (ø)
sbi/inference/trainers/npe/mnpe.py	97.14% <100.00%> (+7.14%)	⬆️
sbi/inference/trainers/npe/npe_base.py	93.38% <ø> (ø)
sbi/inference/trainers/nre/bnre.py	97.29% <100.00%> (+0.07%)	⬆️
sbi/inference/trainers/nre/nre_a.py	100.00% <100.00%> (ø)
sbi/inference/trainers/nre/nre_b.py	100.00% <100.00%> (ø)
sbi/inference/trainers/nre/nre_base.py	92.47% <100.00%> (+1.33%)	⬆️
sbi/inference/trainers/nre/nre_c.py	98.00% <100.00%> (+0.04%)	⬆️
sbi/neural_nets/__init__.py	61.90% <100.00%> (+1.90%)	⬆️
... and 3 more

... and 1 file with indirect coverage changes

[janfb]

Thanks @satwiksps, overall very clean extension of the builder pattern to all four NRE trainers. great that you added the role-shape regression test (test_builder_role_shapes).

As discussed in our meeting I made several longer comments that go beyond the NRE code, mostly improvements that belong in the base class, and although this is a bit of scope creep I think this PR is the right place. It will make the the upcoming VF builder and other refactorings much easier.

Let me know if anything is unclear!

[janfb]

minor, but I would read this more naturally if the and condition would just be inside the if-case:

elif isinstance(classifier, _EstimatorBuilderBase):
        if  not isinstance(classifier, RatioEstimatorBuilder):
            raise TypeError(
                "NRE requires a RatioEstimatorBuilder; got "
                f"{type(classifier).__name__}. Use "
                "RatioEstimatorBuilder(model=...)."
            )
        else:
             self._build_neural_net = self._wrap_builder(classifier)

[janfb]

note it's more a matter of taste and this might apply to the other method's if-else paths as well (NPE, NLE, MNLE, ...). so feel free to leave it like it is now.

[janfb]

as discussed I suggest to change this to @dataclass(frozen=True, eq=False, repr=False) (frozen, no equality because of potential tensors in the classes, and repr False for the printing of config objects (see other comment).

This then has to carry to all subclasses as well, i.e.,: DensityEstimatorBuilder, MixedDensityEstimatorBuilder, RatioEstimatorBuilder, and also the legacy ConditionalFlowConfig / ClassifierConfig / MarginalFlowConfig.

[janfb]

for the docs: when config classes are frozen, dataclasses.replace(builder, hidden_features=64) is the way to make a tweaked copy of a config.

[janfb]

setting self.extra_kwargs = {} in line 45 will not work in the frozen class, so instead we should create it here via extra_kwargs: dict = field(default_factory=dict) (this also removes the # type: ignore).

and then we can delete the if self.extra_kwargs is None: arm below.

[janfb]

as these will be user-facing eventually, let's rename now to more intuitive, but general names: z_score_x → z_score_input, z_score_y → z_score_condition — here and in the density/mixed builders, so all builders agree.

Note, you don't have to refactor all the downstream build_* functions, you can just define an alias in _build_kwargs():

_BUILD_KWARG_ALIASES = {"z_score_input": "z_score_x", "z_score_condition": "z_score_y"}
# in _build_kwargs():
d = {_BUILD_KWARG_ALIASES.get(f.name, f.name): getattr(self, f.name) for f in fields(self) ...}

[janfb]

I suggest to lift these post init checks to the base class level and generalize them across builders: Right now only model is checked and this NRE specific check should stay here , but sth like RatioEstimatorBuilder(z_score_input="typo") does not fail and only fails deep inside .train(). The fix: A generic loop in _EstimatorBuilderBase.__post_init__ gives every builder (this one, density/mixed, and VF next week with its 5 extra Literal fields) construction-time. So, in the base __post_init__ sth like (please double check):

# helper function to collect the corresponding literals (if any). 
def _literal_values(tp) -> frozenset:  # unwraps Optional[Literal[...]]
    if get_origin(tp) is Literal:
        return frozenset(get_args(tp))
    out = frozenset()
    for a in get_args(tp):
        out |= _literal_values(a)
    return out

# in base __post_init__:
for f in fields(self):
    allowed = _literal_values(f.type)
    val = getattr(self, f.name)
    if allowed and val is not None and val not in allowed:
        raise ValueError(f"Invalid value {val!r} for `{f.name}`. Must be one of {sorted(map(str, allowed))} or None.")

[janfb]

Note, I think the model-specific "Unknown model" check (lines 488-492) should probably happen before super().__post_init__() so its friendlier message wins over the generic one and is raised first.

[janfb]

And I have a second general point here with implications for the other builders as well: The NRE builder offers all 9 fields to every model, but linear uses only 4 of them, mlp 6, resnet 8. At the moment, RatioEstimatorBuilder(model="linear", num_blocks=99) silently drops num_blocks (linear's build_* has **kwargs), while model="mlp" with the same field crashes late at build (mlp has no **kwargs), so, same mistake, two different silent failures.
Solution: we check applicability at construction, deriving the allowed set from the build function's signature (so it can never drift from the implementation).

And we can even place the core functions for this into the base class so it can be used by the other builders as well. This was suggested by Claude, so please double check. In the base class, this could look like:

# --- base class: the reusable mechanism ---
@staticmethod
@lru_cache(maxsize=None)
def _fn_param_names(build_fn) -> frozenset:
    return frozenset(
      p.name for p in inspect.signature(build_fn).parameters.values()
        if p.kind is not inspect.Parameter.VAR_KEYWORD and p.name not in ("batch_x", "batch_y"))
def _reject_inapplicable_fields(self, build_fn, *, discriminator, always_ok=frozenset()):
    allowed = self._fn_param_names(build_fn) | always_ok
    set_fields = {f.name for f in fields(self)
                  if f.name not in ("extra_kwargs", discriminator) and getattr(self, f.name) is not None}
    bad = {n for n in set_fields
           if n not in allowed and _BUILD_KWARG_ALIASES.get(n, n) not in allowed}
    if bad:
        raise ValueError(f"Field(s) {sorted(bad)} are not used by {discriminator}="
                         f"{getattr(self, discriminator)!r} and would be silently ignored. "
                         "To forward library-specific options, use `extra_kwargs`.")

and then the subclasses (NRE, NPE, NLE, mixed) post_init decides whether and how to check:

# --- each dispatching builder's __post_init__ opts in ---
def __post_init__(self):
    if self.model not in _VALID_CLASSIFIER_MODELS:   # model-specific msg, first
        raise ValueError(...)
    super().__post_init__()                          # (a) generic Literal check, base
    self._reject_inapplicable_fields(                # (b) opt-in
        _classifier_build_fns()[self.model],         # the shared helper that gives all valid classifier models (see comment below). 
        discriminator="model")                       # Note, this would need to be "continuous_model" for the mixed builder

The discriminator= arg is what lets one base helper serve both model (Density/Ratio) and continuous_model (Mixed) without hard-coding a field name.

[janfb]

the type here could be narrower, e.g., nn.Module as this should always be a NN. At some point we considered allowing Callables that return a nn.Module, but that path is actually not implemented, so no need to type-support it for now. better be more narrow.

this also applies to the other *Builder subclasses.

I also checked those classes for possibly more narrower types and there are a couple of candidates (please check as well), e.g.,:

• MixedDensityEstimatorBuilder.num_categories_per_variable (:386): downstream is Optional[Tensor] → narrow from Any to Optional[Tensor].
• RatioEstimatorBuilder.norm_layer (:484, already Optional[Callable]): downstream is Callable[[int], nn.Module] (default nn.LayerNorm) → tighten to
  Optional[Callable[[int], nn.Module]].

[janfb]

When this gets changed to repr=False from the comment above, you would need to add one custom __repr__ on the base so configs read the way users wrote them. At the moment, repr(RatioEstimatorBuilder(model="resnet", hidden_features=64)) prints all 11 fields starting with the internal extra_kwargs={} and nine Nones, potentially hiding the two the user set. A base __repr__ that leads with model and omits unset/default fields yields RatioEstimatorBuilder(model='resnet', hidden_features=64). (Lives in the base, so all builders benefit.)

[janfb]

When building the repr note that it's important to

• Always show the discriminator (model / continuous_model), even when it equals its default, otherwise DensityEstimatorBuilder(model="maf") (maf is the default) reprs as DensityEstimatorBuilder(), hiding the one field that matters.
• Skip empty extra_kwargs. It uses default_factory, so field.default is MISSING, not {}; a naive value == field.default check won't drop it and you get a stray extra_kwargs={}.

[janfb]

Sorry, I have another comment with implications beyond the initial scope of this PR, but it's better to fix this now: This inline model → build_fn dict duplicates code that already lives in the factories, and the applicability check from the comment above needs the same map, so this calls for a more general helper function, e.g., the _classifier_build_fns that returns the mapping. This applies to the other builders as well, but the Mixed build actually already does it: it delegates to mixed_nets.model_builders (no inline dict) + already has a sync test (PR 4). So, for the applicability check above, it just needs to pass mixed_nets.model_builders[continuous_model].

Howver, for the Ratio* and the Density* builders, we need a helper that just takes the existing builders = {...} dict lifted out of build() into a module-level function, with the import kept inside (mirroring how build() already imports classifier lazily to avoid a circular import):

def _classifier_build_fns() -> dict:
    from sbi.neural_nets.net_builders.classifier import (
        build_linear_classifier,
        build_mlp_classifier,
        build_resnet_classifier,
    )
    return {
        "linear": build_linear_classifier,
        "mlp": build_mlp_classifier,
        "resnet": build_resnet_classifier,
    }

build() then does _classifier_build_fns()[self.model] instead of defining the dict inline. (and analog for _density_build_fns().

Note, this could be made even more general because we have more duplications in the factory.py functions, but this goes beyond this PR. We should fix it later in PR8 or so.

For now, I would just add a sync test per family as a double check so the copies can't drift before then: assert each Literal valid-set equals its build-fn helper's keys, mirroring PR 4's test_valid_continuous_models_match_builders.

[janfb]

Good that NRE strings now warn (they were silent before). I would add the import path so the warning shows the exact line to type, many users won't know where RatioEstimatorBuilder lives:

"Use RatioEstimatorBuilder(model=...) instead, e.g. "
"`from sbi.neural_nets import RatioEstimatorBuilder`."