Fix NaN padding test: remove xfail decorator

[DhyeyJoshi39]

Removed pytest xfail decorator from test_npe_with_iid_embedding_varying_num_trials.

PermutationInvariantEmbedding intentionally uses NaN padding for varying trial counts.
PR #1701's NaN check conflicts with this legitimate use case.

Test passes once check_finite_x=False kwarg is available.

Fixes #1717
Refs: #1701

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 88.51%. Comparing base (649f5d3) to head (2fed695).

Additional details and impacted files

@@           Coverage Diff           @@
##             main    #1735   +/-   ##
=======================================
  Coverage   88.51%   88.51%           
=======================================
  Files         137      137           
  Lines       11527    11527           
=======================================
  Hits        10203    10203           
  Misses       1324     1324           

Flag	Coverage Δ
fast	84.69% <ø> (?)

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

[janfb]

Thank you for making a PR @DhyeyJoshi39. You removed the xfail flag, but I cannot see any implementation of the check_finite_x kwarg.
Did you plan to implement this at a later stage? If so, please mark you PR as draft for now. Thank you! 🙏

[DhyeyJoshi39]

Hi @janfb ,

Thank you for the quick review! I'd like to propose implementing the check_finite_x parameter upstream to resolve this properly.

Problem Statement
PR #1701 added strict isfinite() check to catch invalid data (good!). However, PermutationInvariantEmbedding intentionally pads shorter trial lists with NaNs for varying-length observations. This creates a conflict for legitimate use cases.

Proposed Solution
Add optional check_finite_x parameter (default=True) to:

NPE.build_posterior(check_finite_x=True)

Posterior.sample(check_finite_x=True)

Implementation Sketch:
python

In NPE.build_posterior()

def build_posterior(self, check_finite_x=True, ...):
if check_finite_x:
assert torch.isfinite(x_train).all(), "x contains NaN/Inf!"
# ... rest of logic

In Posterior.sample()

def sample(self, sample_shape, x, check_finite_x=True, ...):
"""
Args:
check_finite_x: If True, raises error if x contains NaN/Inf.
Set False for intentional NaN padding in embeddings like
PermutationInvariantEmbedding.
"""
if check_finite_x:
assert torch.isfinite(x).all(), "x contains NaN/Inf!"
# ... sampling logic
Benefits
Backward compatible - Default True preserves PR #1701 safety
Explicit intent - Flag documents why NaNs are OK
Flexible - Supports PermutationInvariantEmbedding + other special cases
Tested - Unblocks test_npe_with_iid_embedding_varying_num_trials

Questions for Guidance
Should this be in build_posterior(), sample(), or both?

Should we add validation docs explaining when to set False?

Any other files affected by PR #1701 NaN check?

I'm happy to implement this if the approach looks good to you!

Looking forward to your thoughts.

Thanks!

Closes #1717
Related: #1701

[janfb]

Hi @DhyeyJoshi39, thanks for following up on this! Your proposed approach with check_finite_x is the right direction.

A few housekeeping items first:

1. Branch name: Your feature branch is called main, which makes it impossible for us to check out locally for testing. Could you rename it to something like fix-nan-padding or check-finite-x? (See https://sbi.readthedocs.io/en/latest/contributing.html#contribution-workflow)
2. Pre-commit: The CI shows ruff/hooks failing. Please run pre-commit install and then pre-commit run --all-files locally before pushing (see https://sbi.readthedocs.io/en/latest/contributing.html#style-conventions-and-testing).

---

Now for the implementation. The test change alone won't work because check_finite_x doesn't exist yet. Here's what's needed:

Background

There's an asymmetry right now:

• Training: append_simulations(..., exclude_invalid_x=False) allows NaN-padded data via handle_invalid_x() in sbi/utils/sbiutils.py
• Inference: process_x() unconditionally calls assert_all_finite(), which was added in assert finite x_o #1701

There's even a TODO in handle_invalid_x() acknowledging this use case:

# TODO: add option to allow for NaNs in certain dimensions, e.g., to encode varying
# numbers of trials.

Implementation

The NaN padding use case is specific to NPE with PermutationInvariantEmbedding for trial-based data.

1. process_x() in sbi/utils/user_input_checks.py:

def process_x(
    x: Array,
    x_event_shape: Optional[torch.Size] = None,
    check_finite: bool = True,
) -> Tensor:
    ...
    x = atleast_2d(torch.as_tensor(x, dtype=float32))
    if check_finite:
        assert_all_finite(x, "Observed data x_o contains Nans or Infs.")
    ...

2. NeuralPosterior.__init__() in sbi/inference/posteriors/base_posterior.py:

• Add check_finite_x: bool = True parameter
• Store as self._check_finite_x

3. NeuralPosterior._x_else_default_x() and set_default_x() (same file):

• Pass check_finite=self._check_finite_x to process_x()

4. DirectPosterior.__init__() in sbi/inference/posteriors/direct_posterior.py:

• Add check_finite_x: bool = True parameter
• Pass to super().__init__(..., check_finite_x=check_finite_x)

5. PosteriorEstimatorTrainer.build_posterior() in sbi/inference/trainers/npe/npe_base.py:

• Add check_finite_x: bool = True parameter
• Pass to DirectPosterior(..., check_finite_x=check_finite_x)

That's 5 functions to modify. The other posterior types inherit _check_finite_x from NeuralPosterior but will default to True, so they behave exactly as before.

Let me know if you have questions about any of this!