nn module as base class for density estimator

[manuelgloeckler]

Density Estimators inherit nn.Module.
Also added sample_and_log_prob, with a naive implementation that can be easily overriden (following Zuko API).

Limitations:

• net must be a nn.Module (which is fine and consistent with the type hints).

[codecov]

Codecov Report

Attention: Patch coverage is 90.38462% with 5 lines in your changes are missing coverage. Please review.

Project coverage is 76.17%. Comparing base (c3c2b6e) to head (082dd8d).

Files	Patch %	Lines
sbi/neural_nets/density_estimators/base.py	66.66%	5 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #965      +/-   ##
==========================================
+ Coverage   76.08%   76.17%   +0.08%     
==========================================
  Files          83       83              
  Lines        6361     6406      +45     
==========================================
+ Hits         4840     4880      +40     
- Misses       1521     1526       +5     

Flag	Coverage Δ
unittests	76.17% <90.38%> (+0.08%)	⬆️

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

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[manuelgloeckler]

A quick question:

The sample for NFlowsFlow with e.g. sample shape (n,) will given a batched condition (b, d) return a tensor of shape (10, b*d). Is this intended? To be consistent with torch we should return (10, b,d), no?

[gmoss13]

A quick question:

The sample for NFlowsFlow with e.g. sample shape (n,) will given a batched condition (b, d) return a tensor of shape (10, b*d). Is this intended? To be consistent with torch we should return (10, b,d), no?

I agree that it should be batched! But I believe that nflows in this case would return the condition batch dimension first in such a case, i.e. (b, n ,d2) where d2 is the dimension of the inputs (and d is the dimension of the condition). I would be in favour of sticking to this convention.

[manuelgloeckler]

I see and I agree, I will add (b, n, d2) support.

[gmoss13]

Thanks for the PR! I left a few small comments, happy to discuss.

[janfb]

thanks a lot for adding this!
A added a couple of comments and questions.

[manuelgloeckler]

Sampling should now work with arbitrary context ndim and added a bunch of tests.

One question @gmoss13 @janfb, the log_prob as it was implemented currently only works if input.shape[0] == condition.shape[0] (by nflows).

This is quite restrictive and always requires manually reshaping and repeating. This is fine for the loss, but we might want to make this more general for the log_prob method.

• If we have an input of shape (n, d1) and condition (d2,) -> (n,) (automate repeat)
• If we have an input of shape (n, d1) and condition (m,d2) -> (m,n) ? (this is in conflict with previous behavior for m=n)

I think to mirror PyTorch behavior one would adapt "broadcasting" semantics to resolve the second point i.e.:

• (n,1,d1), (m, d2) -> (n,m)
• (n, d1), (m,d2) -> error (n != m) or (n,) for (m == n)

[manuelgloeckler]

Alright, I added to the DensityEstimator

• x_shape to density estimator + a basic check if inputs are consistent with this attribute

To NFlowFlows:

• log_prob through broadcasting (a bit ugly due to nflow limitaitons).
• Using x_shape to squeeze/unsqueeze/reshape the right dimensions.
• Extended test cases.

Have to add a bit more documentation i.e. what we expect each density estimator can do.

[janfb]

Great work!
Added a couple of comment about variable names.

[janfb]

Thanks! almost ready I think.
Just needs clarification for the broadcasting and some more docstrings 🙄 sorry 😄

[janfb]

thanks for the clarifications!
Feel free to re-assign me for a final review once everything is addressed.

[manuelgloeckler]

Thanks for your comments. I tried to fix all left-over inconsistencies now.

[janfb]

Looks good, thanks!