Use newest version of pyright, ignored a lot of errors, identified some bugs

[Baschdl]

I worked through the errors pyright raises on the newest version. I used very small commits to be able to cherry pick if we only want to go with some of the decisions I made. You can rerun stuff with pyright locally (you can also find the errors in the CI run).

I believe that the remaining errors are actual errors for which we don't have unit tests currently (I tagged the people who seem to be responsible for those parts of the code):

• @gmoss13 Error in analysis/conditional_density.py:
  

  sbi/sbi/analysis/conditional_density.py

  Line 209 in 1d25046

  logits, means, precfs, _ = extract_and_transform_mog(net=mdn.net, context=x_o)

  
  Here you assume that the network of your new DensityEstimator has the same interface as Flow. As an example, I would assume that the general nn.Module doesn't necessarily has the _distribution attribute we access later on:
  

  sbi/sbi/utils/conditional_density_utils.py

  Line 179 in 1d25046

  dist = net._distribution

• @gmoss13 Error in neural_nets/mnle.py:
  

  sbi/sbi/neural_nets/mnle.py

  Line 96 in c347d5c

  continuous_net=cont_nle,

  
  This error currently only says that the __init__ of MixedDensityEstimator doesn't expect aNFlowsFlow:
  

  sbi/sbi/neural_nets/mnle.py

  Line 217 in c347d5c

  continuous_net: flows.Flow,

  
  The obvious thing would be to add it as a possible type but the problem is that the used calls to sample, e.g.
  

  sbi/sbi/neural_nets/mnle.py

  Line 265 in c347d5c

  continuous_x = self.continuous_net.sample(

  
  are not compatible to your DensityEstimator.
• @janfb Errors in utils/user_input_checks.py and user_input_checks_utils.py:
  

  sbi/sbi/utils/user_input_checks_utils.py

  Line 151 in c347d5c

  return self.prior_scipy.var

  
  You assume that univariate and multivariate random variables have the same interface. I've already added a failing test in b455a7f
• Maybe it's worth having another look into e9d8dff as it was unclear to me if LazyTransform and AffineTransform | Unconditional are compatible.

Feel free to fix things directly in this PR or create issues for those problems. We won't be able to merge this PR before fixing them.

Fixes #847

Checklist

• I agree with re-licensing my contribution from AGPLv3 to Apache-2.0.
• I performed linting and formatting as described in the contribution
  guidelines

[gmoss13]

Thanks @Baschdl, great work! Some comments:

@gmoss13 Error in analysis/conditional_density.py:

sbi/sbi/analysis/conditional_density.py

Line 209 in 1d25046
logits, means, precfs, _ = extract_and_transform_mog(net=mdn.net, context=x_o)

Here you assume that the network of your new DensityEstimator has the same interface as Flow. As an example, I would assume that the general nn.Module doesn't necessarily has the _distribution attribute we access later on:

sbi/sbi/utils/conditional_density_utils.py

Line 179 in 1d25046
dist = net._distribution

These functions only apply to mixture density networks - and should only be called when the density estimator is an MDN. Our implementation of MDN is using normflows, and so this function should only be called when the DensityEstimator is specifically an instance of NFlowsFlow. I can jump in here and explicitly change the type for the mdn argument, and clarify this point in the docstring.

@gmoss13 Error in neural_nets/mnle.py:

sbi/sbi/neural_nets/mnle.py

Line 96 in c347d5c
continuous_net=cont_nle,

This error currently only says that the init of MixedDensityEstimator doesn't expect aNFlowsFlow:

sbi/sbi/neural_nets/mnle.py

Line 217 in c347d5c
continuous_net: flows.Flow,

The obvious thing would be to add it as a possible type but the problem is that the used calls to sample, e.g.

sbi/sbi/neural_nets/mnle.py

Line 265 in c347d5c
continuous_x = self.continuous_net.sample(

are not compatible to your DensityEstimator.

Yes - this should resolved by updating the MixedDensityEstimator to the new DensityEstimator interface (see #968) - so I suggest we wait for #968 to be resolved before merging this @coschroeder .

[michaeldeistler]

Infinite karma for this! A few requests below, but feel free to merge then.

[michaeldeistler]

+1 for this, thank you!

[michaeldeistler]

why remove this?

[Baschdl]

@janfb's assessment was that no one is using this wrapper, so instead of trying to fix this issue he suggested to deprecate it.

[Baschdl]

I'll merge it for now but we can still bring it back. Let me know.

[janfb]

great effort! 👏
Karma +1000