poutine.reparam(model,-) breaks AutoGuide initialization

[fritzo]

poutine.reparam(model, -) currently replaces targeted model sites with deterministic sites and introduces auxiliary latent variables upstream. While this site swapping is mostly compatible with AutoGuides, it does break autoguide initialization, since init_loc_fn(-) initializes based on the original site name.

Despite this breakage, Reparam objects generally know how to transform between downstream original latents and upstream auxiliary latents. Thus there is an opportunity to refactor AutoGuide initialization and/or Reparam.__call__() so that autoguides are more agnostic to reparameterization.

As proof of concept, consider adding inside each reparametrizer a new non-traced site that is used only for initialization:

  class MyReparam(Reparam):
      def __call__(self, name, fn, value):
            ...
+           if initializing:  # read from some global state?
+               value = pyro.sample(name, fn, infer={"_do_not_trace": True})
+               value_t = my_transform(value)
+               with poutine.condition(data={name + "_transformed": transformed_value}):
+                   return self(name, fn, value)
          ...            

[fehiepsi]

cc @martinjankowiak - who also concerned about this

[fritzo]

I believe this was fixed by #2876