Add Bean Machine VI to PPL Bench

examples/robust_regression.json

@@ -59,7 +59,21 @@
         }
       },
       "legend": {
-        "color": "purple"
+        "color": "purple",
+        "name": "beanmachine-NUTS"
+      }
+    },
+    {
+      "name": "beanmachine",
+      "inference": {
+        "class": "inference.VI",
+        "infer_args": {
+          "algorithm": "ADVI"
+        }
+      },
+      "legend": {
+        "color": "blue",
+        "name": "beanmachine-ADVI"
       }
     },
     {
@@ -68,7 +82,7 @@
         "class": "inference.MCMC",
         "infer_args": {"algorithm": "NUTS"}
       },
-      "legend": {"color": "orange", "name": "numpyro"}
+      "legend": {"color": "orange", "name": "numpyro-NUTS"}
     }
   ],
   "save_samples": true,

pplbench/ppls/beanmachine/inference.py

@@ -9,6 +9,8 @@
 import torch
 import xarray as xr
 from beanmachine.ppl import inference
+from beanmachine.ppl.inference.monte_carlo_samples import MonteCarloSamples
+from beanmachine.ppl.inference.vi.variational_world import VariationalWorld
 
 from ..base_ppl_impl import BasePPLImplementation
 from ..base_ppl_inference import BasePPLInference
@@ -51,3 +53,42 @@ def infer(
             num_chains=1,
         )
         return self.impl.extract_data_from_bm(samples)
+
+
+class VI(BaseBMInference):
+    def infer(
+        self,
+        data: xr.Dataset,
+        iterations: int,
+        num_warmup: int,
+        seed: int,
+        algorithm: str = "ADVI",
+        **infer_args,
+    ) -> xr.Dataset:
+        bm.seed(seed)
+        inference_cls = getattr(inference.vi, algorithm)
+        if not issubclass(inference_cls, inference.vi.autoguide.AutoGuideVI):
+            raise ValueError("Only autoguide methods are supported in PPL Bench")
+
+        vi_world = inference_cls(
+            queries=self.impl.get_queries(),
+            observations=self.impl.data_to_observations(data),
+            **infer_args,
+        ).infer(num_steps=iterations)
+
+        samples = self._draw_mc_samples_from_vi_world(vi_world, iterations)
+
+        return self.impl.extract_data_from_bm(samples)
+
+    def _draw_mc_samples_from_vi_world(
+        self, vi_world: VariationalWorld, num_samples: int
+    ) -> MonteCarloSamples:
+        samples = {}
+        for query in self.impl.get_queries():
+            samples[query] = (
+                vi_world.get_guide_distribution(query)
+                .sample((num_samples,))
+                .clone()
+                .detach()
+            )
+        return MonteCarloSamples([samples])