Fix: Make NPSE picklable (#1679)

Author: michaeldeistler

sbi/neural_nets/estimators/score_estimator.py

@@ -399,21 +399,28 @@ def _set_weight_fn(self, weight_fn: Union[str, Callable]):
                 - a custom function that returns a Callable.
         """
         if weight_fn == "identity":
-            self.weight_fn = lambda times: 1
+            self.weight_fn = self._identity_weight_fn
         elif weight_fn == "max_likelihood":
-            self.weight_fn = (
-                lambda times: self.diffusion_fn(
-                    torch.ones((1,), device=times.device), times
-                )
-                ** 2
-            )
+            self.weight_fn = self._max_likelihood_weight_fn
         elif weight_fn == "variance":
-            self.weight_fn = lambda times: self.std_fn(times) ** 2
+            self.weight_fn = self._variance_weight_fn
         elif callable(weight_fn):
             self.weight_fn = weight_fn
         else:
             raise ValueError(f"Weight function {weight_fn} not recognized.")
 
+    def _identity_weight_fn(self, times):
+        """Return ones for any time t."""
+        return 1
+
+    def _max_likelihood_weight_fn(self, times):
+        """Return weights proportional to the diffusion function."""
+        return self.diffusion_fn(torch.ones((1,), device=times.device), times) ** 2
+
+    def _variance_weight_fn(self, times):
+        """Return weights as the variance."""
+        return self.std_fn(times) ** 2
+
     def ode_fn(self, input: Tensor, condition: Tensor, times: Tensor) -> Tensor:
         """ODE flow function of the score estimator.
 

sbi/samplers/ode_solvers/zuko_ode.py

@@ -5,6 +5,8 @@
 Zuko ODE solver.
 """
 
+import functools
+
 import torch.nn as nn
 from torch import Tensor
 from torch.distributions import Distribution
@@ -81,8 +83,9 @@ def get_distribution(self, condition: Tensor, **kwargs) -> Distribution:
             The distribution object with `log_prob` and
             `sample` methods that wraps the ODE solver.
         """
+        partial_f = functools.partial(self._f_condition_last, condition=condition)
         transform = FreeFormJacobianTransform(
-            f=lambda t, input: self.f(input, condition, t),
+            f=partial_f,
             t0=condition.new_tensor(self.t_min),
             t1=condition.new_tensor(self.t_max),
             phi=(condition, *self.net.parameters()),
@@ -93,3 +96,10 @@ def get_distribution(self, condition: Tensor, **kwargs) -> Distribution:
             transform=transform,
             base=DiagNormal(self.mean_base, self.std_base).expand(condition.shape[:-1]),
         )
+
+    def _f_condition_last(self, t, input, condition):
+        """Return the ODE function with argument-order changed: condition is last.
+
+        This is a helper to build a partial function in `get_distribution`, which
+        requires that `condition` is passed as last element."""
+        return self.f(input, condition, t)

tests/save_and_load_test.py

@@ -7,12 +7,13 @@
 import torch
 
 from sbi import utils as utils
-from sbi.inference import NLE, NPE, NRE
+from sbi.inference import FMPE, NLE, NPE, NPSE, NRE
 from sbi.inference.posteriors.posterior_parameters import (
     DirectPosteriorParameters,
     MCMCPosteriorParameters,
     RejectionPosteriorParameters,
     VIPosteriorParameters,
+    VectorFieldPosteriorParameters,
 )
 from sbi.inference.posteriors.vi_posterior import VIPosterior
 
@@ -21,6 +22,8 @@
     "inference_method, posterior_parameters",
     (
         (NPE, DirectPosteriorParameters),
+        (NPSE, VectorFieldPosteriorParameters),
+        (FMPE, VectorFieldPosteriorParameters),
         pytest.param(NLE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
         pytest.param(NRE, MCMCPosteriorParameters, marks=pytest.mark.mcmc),
         pytest.param(NRE, VIPosteriorParameters, marks=pytest.mark.mcmc),