fix typing and score gradient bug

sbi-dev · Aug 20, 2024 · 5071e43 · 5071e43
1 parent a03b711
commit 5071e43
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 10 deletions.
diff --git a/sbi/samplers/score/kernels/base.py b/sbi/samplers/score/kernels/base.py
@@ -2,9 +2,7 @@
 
 from torch import Tensor
 
-from sbi.inference.potentials.score_based_potential import (
-    PosteriorScoreBasedPotentialGradient,
-)
+from sbi.inference.potentials.score_based_potential import PosteriorScoreBasedPotential
 
 
 class State:
@@ -14,11 +12,13 @@ def __init__(self, input: Tensor, time: Tensor) -> None:
 
 
 class Kernel:
-    def __init__(self, score_fn: PosteriorScoreBasedPotentialGradient) -> None:
+    def __init__(self, score_fn: PosteriorScoreBasedPotential) -> None:
         self.score_fn = score_fn
 
     @abstractmethod
-    def __call__(self, state: State, time: Tensor) -> Tensor:
+    def __call__(
+        self, state: State, time: Tensor, track_gradients: bool = True
+    ) -> Tensor:
         pass
 
     def __repr__(self):

diff --git a/sbi/samplers/score/kernels/euler_maruyama.py b/sbi/samplers/score/kernels/euler_maruyama.py
@@ -3,14 +3,12 @@
 import torch
 from torch import Tensor
 
-from sbi.inference.potentials.score_based_potential import (
-    PosteriorScoreBasedPotentialGradient,
-)
+from sbi.inference.potentials.score_based_potential import PosteriorScoreBasedPotential
 from sbi.samplers.score.kernels.base import Kernel, State
 
 
 class EulerMaruyama(Kernel):
-    def __init__(self, score_fn: PosteriorScoreBasedPotentialGradient, eta=1.0) -> None:
+    def __init__(self, score_fn: PosteriorScoreBasedPotential, eta=1.0) -> None:
         self.score_fn = score_fn
         self.drift_forward = score_fn.score_estimator.drift_fn
         self.diffusion_forward = score_fn.score_estimator.diffusion_fn
@@ -23,7 +21,8 @@ def __call__(self, state: State, time: Tensor) -> State:
         delta_t = time - time_old
         f = self.drift_forward(input_old, time_old)
         g = self.eta * self.diffusion_forward(input_old, time_old)
-        score = self.score_fn(input_old, time_old)
+        # Call the gradient of the score.
+        score = self.score_fn.gradient(input_old, time_old)
         f_backward = f - (1 + self.eta**2) / 2 * g**2 * score
 
         new_input = (