Fixing bug on merge

Author: manuelgloeckler

sbi/inference/posteriors/vector_field_posterior.py

@@ -404,7 +404,13 @@ def log_prob(
             `(len(θ),)`-shaped log posterior probability $\log p(\theta|x)$ for θ in the
             support of the prior, -∞ (corresponding to 0 probability) outside.
         """
-        self.potential_fn.set_x(self._x_else_default_x(x), **(ode_kwargs or {}))
+        x = self._x_else_default_x(x)
+        x = reshape_to_batch_event(x, self.vector_field_estimator.condition_shape)
+        is_iid = x.shape[0] > 1
+        self.potential_fn.set_x(
+            x,
+            x_is_iid=is_iid,
+        )
 
         theta = ensure_theta_batched(torch.as_tensor(theta))
         return self.potential_fn(

sbi/inference/potentials/vector_field_potential.py

@@ -115,8 +115,6 @@ def set_x(
         super().set_x(x_o, x_is_iid)
         self.iid_method = iid_method or self.iid_method
         self.iid_params = iid_params
-        # NOTE: Once IID potential evaluation is supported. This needs to be adapted.
-        # See #1450.
         if not x_is_iid and (self._x_o is not None):
             self.flow = self.rebuild_flow(**ode_kwargs)
         elif self._x_o is not None:
@@ -163,6 +161,8 @@ def __call__(
                     ),
                     dim=0,
                 )
+                # Apply the adjustment for iid observations i.e. we have to subtract
+                # (n-1) times the log prior.
                 log_probs = iid_posteriors_prob - (n - 1) * self.prior.log_prob(
                     theta_density_estimator
                 ).squeeze(-1)