Replace some torch.vmap usage with a hand-vectorized BatchedNMCLogPredictiveLikelihood

chirho/observational/handlers/condition.py

@@ -1,10 +1,10 @@
-from typing import Callable, Generic, Hashable, Mapping, TypeVar, Union
+from typing import Callable, Generic, Mapping, TypeVar, Union
 
 import pyro
 import torch
 
 from chirho.observational.internals import ObserveNameMessenger
-from chirho.observational.ops import AtomicObservation, observe
+from chirho.observational.ops import Observation, observe
 
 T = TypeVar("T")
 R = Union[float, torch.Tensor]
@@ -62,7 +62,9 @@ class Observations(Generic[T], ObserveNameMessenger):
     a richer set of observational data types and enables counterfactual inference.
     """
 
-    def __init__(self, data: Mapping[Hashable, AtomicObservation[T]]):
+    data: Mapping[str, Observation[T]]
+
+    def __init__(self, data: Mapping[str, Observation[T]]):
         self.data = data
         super().__init__()
 

chirho/robust/internals/linearize.py

@@ -5,7 +5,7 @@
 import torch
 from typing_extensions import Concatenate, ParamSpec
 
-from chirho.robust.internals.predictive import NMCLogPredictiveLikelihood
+from chirho.robust.internals.predictive import BatchedNMCLogPredictiveLikelihood
 from chirho.robust.internals.utils import (
     ParamDict,
     make_flatten_unflatten,
@@ -92,23 +92,17 @@ def f_Ax_flat(v: torch.Tensor) -> torch.Tensor:
 
 
 def make_empirical_fisher_vp(
-    func_log_prob: Callable[Concatenate[ParamDict, Point[T], P], torch.Tensor],
+    batched_func_log_prob: Callable[Concatenate[ParamDict, Point[T], P], torch.Tensor],
     log_prob_params: ParamDict,
     data: Point[T],
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> Callable[[ParamDict], ParamDict]:
-    batched_func_log_prob: Callable[[ParamDict, Point[T]], torch.Tensor] = torch.vmap(
-        lambda p, data: func_log_prob(p, data, *args, **kwargs),
-        in_dims=(None, 0),
-        randomness="different",
-    )
-
     N = data[next(iter(data))].shape[0]  # type: ignore
     mean_vector = 1 / N * torch.ones(N)
 
     def bound_batched_func_log_prob(params: ParamDict) -> torch.Tensor:
-        return batched_func_log_prob(params, data)
+        return batched_func_log_prob(params, data, *args, **kwargs)
 
     def _empirical_fisher_vp(v: ParamDict) -> ParamDict:
         def jvp_fn(log_prob_params: ParamDict) -> torch.Tensor:
@@ -145,12 +139,11 @@ def linearize(
         num_samples=num_samples_outer,
         parallel=True,
     )
-    predictive_params, func_predictive = make_functional_call(predictive)
 
-    log_prob = NMCLogPredictiveLikelihood(
+    batched_log_prob = BatchedNMCLogPredictiveLikelihood(
         model, guide, num_samples=num_samples_inner, max_plate_nesting=max_plate_nesting
     )
-    log_prob_params, func_log_prob = make_functional_call(log_prob)
+    log_prob_params, batched_func_log_prob = make_functional_call(batched_log_prob)
     log_prob_params_numel: int = sum(p.numel() for p in log_prob_params.values())
     if cg_iters is None:
         cg_iters = log_prob_params_numel
@@ -166,23 +159,18 @@ def _fn(
         **kwargs: P.kwargs,
     ) -> ParamDict:
         with torch.no_grad():
-            data: Point[T] = func_predictive(predictive_params, *args, **kwargs)
+            data: Point[T] = predictive(*args, **kwargs)
             data = {k: data[k] for k in points.keys()}
         fvp = make_empirical_fisher_vp(
-            func_log_prob, log_prob_params, data, *args, **kwargs
+            batched_func_log_prob, log_prob_params, data, *args, **kwargs
         )
         pinned_fvp = reset_rng_state(pyro.util.get_rng_state())(fvp)
         pinned_fvp_batched = torch.func.vmap(
             lambda v: pinned_fvp(v), randomness="different"
         )
-        batched_func_log_prob = torch.vmap(
-            lambda p, data: func_log_prob(p, data, *args, **kwargs),
-            in_dims=(None, 0),
-            randomness="different",
-        )
 
         def bound_batched_func_log_prob(p: ParamDict) -> torch.Tensor:
-            return batched_func_log_prob(p, points)
+            return batched_func_log_prob(p, points, *args, **kwargs)
 
         if pointwise_influence:
             score_fn = torch.func.jacrev(bound_batched_func_log_prob)