Add a Trace_ELBO wrapper to minipyro (pyro-ppl#1794)

examples/minipyro.py

@@ -11,7 +11,7 @@
 import torch
 
 # We use the pyro.generic interface to support dynamic choice of backend.
-from pyro.generic import backend
+from pyro.generic import pyro_backend
 from pyro.generic import distributions as dist
 from pyro.generic import infer, optim, pyro
 
@@ -37,11 +37,12 @@ def guide(data):
 
     # Because the API in minipyro matches that of Pyro proper,
     # training code works with generic Pyro implementations.
-    with backend(args.backend):
+    with pyro_backend(args.backend):
         # Construct an SVI object so we can do variational inference on our
-        # model/guide pair. We work around small differences in elbo interface.
-        elbo = infer.elbo if args.backend == "minipyro" else infer.Trace_ELBO()
-        svi = infer.SVI(model, guide, optim.Adam({"lr": args.learning_rate}), elbo)
+        # model/guide pair.
+        elbo = infer.Trace_ELBO()
+        adam = optim.Adam({"lr": args.learning_rate})
+        svi = infer.SVI(model, guide, adam, elbo)
 
         # Basic training loop
         pyro.get_param_store().clear()

pyro/contrib/minipyro.py

@@ -260,7 +260,7 @@ def step(self, *args, **kwargs):
 # fundamental objective in Variational Inference.
 # See http://pyro.ai/examples/svi_part_i.html for details.
 # This implementation has various limitations (for example it only supports
-# random variablbes with reparameterized samplers), but all the ELBO
+# random variables with reparameterized samplers), but all the ELBO
 # implementations in Pyro share the same basic logic.
 def elbo(model, guide, *args, **kwargs):
     # Run the guide with the arguments passed to SVI.step() and trace the execution,
@@ -289,3 +289,8 @@ def elbo(model, guide, *args, **kwargs):
     # Return (-elbo) since by convention we do gradient descent on a loss and
     # the ELBO is a lower bound that needs to be maximized.
     return -elbo
+
+
+# This is a wrapper for compatibility with full Pyro.
+def Trace_ELBO(*args, **kwargs):
+    return elbo

pyro/generic.py

@@ -30,7 +30,7 @@ def __getattribute__(self, name):
 
 
 @contextmanager
-def backend(*aliases, **new_backends):
+def pyro_backend(*aliases, **new_backends):
     """
     Context manager to set a custom backend for Pyro models.