Add RatioEstimatorBuilder for NRE trainers and base class improvements by satwiksps · Pull Request #1920 · sbi-dev/sbi

satwiksps · 2026-07-04T10:30:11Z

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 · 2026-07-06T14:12:11Z

Codecov Report

❌ Patch coverage is 99.23664% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 87.98%. Comparing base (587f922) to head (2896bdc).
⚠️ Report is 38 commits behind head on gsoc-2026.
✅ All tests successful. No failed tests found.

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

Additional details and impacted files

@@              Coverage Diff              @@
##           gsoc-2026    #1920      +/-   ##
=============================================
- Coverage      87.99%   87.98%   -0.02%     
=============================================
  Files            144      144              
  Lines          13499    13717     +218     
=============================================
+ Hits           11879    12069     +190     
- Misses          1620     1648      +28

Flag	Coverage Δ
fast	`81.69% <99.23%> (?)`

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.05% <ø> (+4.46%)`	⬆️
sbi/inference/trainers/nle/nle_base.py	`97.36% <ø> (ø)`
sbi/inference/trainers/npe/mnpe.py	`97.05% <ø> (+7.05%)`	⬆️
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 2 files 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 · 2026-07-08T15:31:24Z

+    embedding_net_x: Optional[Any] = None
+    embedding_net_y: Optional[Any] = None


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 · 2026-07-08T17:25:34Z

+_VALID_CLASSIFIER_MODELS = frozenset(get_args(CLASSIFIER_MODELS))
+
+
+@dataclass


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.)

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={}.

…idation * Converted all builder classes to use @DataClass(frozen=True, eq=False, repr=False) to ensure immutability. * Added _reject_inapplicable_fields() to fail-fast and throw an error if a user passes hyperparameters that the chosen model's build function would silently ignore. * Implemented a custom __repr__ that hides empty or default fields but always displays the discriminator field (like model or continuous_model). * Replaced the None default for extra_kwargs with ield(default_factory=dict) to handle arbitrary kwargs safely. * Renamed the z-score fields to z_score_input and z_score_condition, while using a _BUILD_KWARG_ALIASES map to maintain compatibility with downstream build functions. * Added _literal_values() to catch typos in Literal fields immediately when the class is constructed. * Tightened the type annotations for fields like embedding_net, um_categories_per_variable, and orm_layer. * Moved the dispatch dictionaries out to module-level _density_build_fns() and _classifier_build_fns() helpers to keep the code clean and prevent circular imports.

janfb

Looking mostly done, great! 👏

Just one missed argument and a couple of smaller improvement comments.

janfb

Looks good! 👏

satwiksps force-pushed the nn-builder-refactor-pr-five branch from d442dce to 8d5def0 Compare July 4, 2026 10:43

satwiksps marked this pull request as ready for review July 6, 2026 14:08

janfb requested changes Jul 8, 2026

View reviewed changes

satwiksps changed the title ~~Add RatioEstimatorBuilder for NRE trainers~~ Add RatioEstimatorBuilder for NRE trainers and base class improvements Jul 20, 2026

satwiksps force-pushed the nn-builder-refactor-pr-five branch from 75f7a74 to a7c12dd Compare July 20, 2026 06:56

satwiksps added 13 commits July 20, 2026 20:07

add RatioEstimatorBuilder for NRE classifiers

cc28767

integrate RatioEstimatorBuilder into RatioEstimatorTrainer

cf33b9c

support RatioEstimatorBuilder with None default

a641096

fix docstrings across all files

8fe4a59

add tests for nre builder integration

058814c

expose the new APIs to public

8b11b83

use nested if case instead of and

da19897

make dataclass completely immutable

d273246

add noqa to fix linting error and some docstring changes

883d65c

add test for new z_score_alias

9618250

add more integration tests for RatioEstimatorBuilder

77ddbeb

fix new naming conventions in tests

8122159

satwiksps force-pushed the nn-builder-refactor-pr-five branch from a7c12dd to 8122159 Compare July 20, 2026 14:37

satwiksps requested a review from janfb July 20, 2026 17:23

janfb requested changes Jul 21, 2026

View reviewed changes

satwiksps added 6 commits July 22, 2026 05:10

revert to using and condition for mnle and mnpe

a450581

add back to docstrings what the classifier is

7f63213

add dropout_probability to prevent ValueError

d35e38a

remove _build_kwargs repetition and _VALID_MIXED_CONTINUOUS_MODELS list

9e1544c

replace _VALID_MIXED_CONTINUOUS_MODELS list in tests

d4aabae

style: format estimator_configs.py with ruff

2896bdc

satwiksps requested a review from janfb July 22, 2026 05:57

janfb approved these changes Jul 22, 2026

View reviewed changes

janfb merged commit 37d0327 into sbi-dev:gsoc-2026 Jul 22, 2026
21 checks passed

janfb mentioned this pull request Jul 22, 2026

release: prepare 0.27.0 #1947

Open

7 tasks

		embedding_net_x: Optional[Any] = None
		embedding_net_y: Optional[Any] = None

		_VALID_CLASSIFIER_MODELS = frozenset(get_args(CLASSIFIER_MODELS))


		@dataclass

Uh oh!

Conversation

satwiksps commented Jul 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Files Changed

AI Usage

Does this close any issues?

Any relevant code examples, logs, or error messages?

Uh oh!

codecov Bot commented Jul 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

janfb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janfb Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

janfb Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

janfb Jul 8, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

janfb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

janfb left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

satwiksps commented Jul 4, 2026 •

edited

Loading

codecov Bot commented Jul 6, 2026 •

edited

Loading