automated robust inference clean up

chirho/robust/ops.py

@@ -19,8 +19,8 @@ def influence_fn(
     model: Callable[P, Any],
     guide: Callable[P, Any],
     functional: Optional[Functional[P, S]] = None,
-    **linearize_kwargs
-) -> Callable[Concatenate[Point[T], P], S]:
+    **linearize_kwargs,
+) -> Callable[Concatenate[Point[T], bool, P], S]:
     from chirho.robust.internals.linearize import linearize
     from chirho.robust.internals.predictive import PredictiveFunctional
     from chirho.robust.internals.utils import make_functional_call
@@ -43,7 +43,7 @@ def _fn(
         points: Point[T],
         pointwise_influence: bool = False,
         *args: P.args,
-        **kwargs: P.kwargs
+        **kwargs: P.kwargs,
     ) -> S:
         param_eif = linearized(
             points, pointwise_influence=pointwise_influence, *args, **kwargs
@@ -53,6 +53,21 @@ def _fn(
                 lambda p: func_target(p, *args, **kwargs), (target_params,), (d,)
             )[1],
             in_dims=0,
+            randomness="different",
         )(param_eif)
 
     return _fn
+
+
+def one_step_correction(
+    model: Callable[P, Any],
+    guide: Callable[P, Any],
+    test_data: Point[T],
+    functional: Optional[Functional[P, S]] = None,
+    **influence_kwargs,
+) -> Callable[P, S]:
+    def _one_step(*args, **kwargs) -> S:
+        eif_fn = influence_fn(model, guide, functional, **influence_kwargs)
+        return eif_fn(test_data, pointwise_influence=False, *args, **kwargs)
+
+    return _one_step