Merging My First PR for GSoC!

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!