Delay log density computation in pyro.factor

In several places (e.g. NeuTra reparam) that we used pyro.factor or dist.Delta, log_density is computed explicitly. This computation is not needed, e.g. when making predictions. It would be nice to add a mechanism to defer this computation, only until some inference algorithms that need the log probability.

Proposed solution

Allow log_factor to be a callable in pyro.factor. In dist.Unit.log_prob, if log_factor is a callable, we will call it to get the required value (similar to the behavior in pyro.param).

[fritzo]

Looking at more real-world models, I see there is often other code surrounding pyro.factor, including plates and poutines:

def model():
    ...
    # This whole section could be omitted during prediction:
    tree_scale = pyro.sample("tree_scale", dist.LogNormal(0, 5))
    with pyro.plate("edges", len(edges), dim=-1):
        u, v = log_rate[edges].unbind(-1)
        pyro.sample(
            "rate_change",
            dist.Laplace(0, tree_scale),
            obs=u - v,
        )
    ...

One advantage of a check like is_masked() is that we could use it around entire sections of model code. In particular plates and poutines would not be able to be embedded inside a lambda passed to pyro.factor.

[fritzo]

Here’s a possible solution that avoids the need for non-tensors in distributions (which would complicate typing):

What if we use poutine.mask(mask=False) to control whether various computations are needed? During prediction we can install a mask(False) handler to disable computation of log density. A brute force way to consult the mask would be to add a primitive is_mask_false() or something. Or instead of a new poutine, we could in NeuTra add tracing logic like this:

- z_unconstrained = pyro.sample("{}_shared_latent".format(name),
-                               self.guide.get_base_dist().mask(False))
+ with poutine.trace() as trace:
+     z_name = f"{name}_shared_latent"
+     z_unconstrained = pyro.sample(z_name, self.guide.get_base_dist().mask(False))
+ z_mask = trace.nodes[z_name]["mask"]
  
  # Differentiably transform.
  x_unconstrained = self.transform(z_unconstrained)
- log_density = self.transform.log_abs_det_jacobian(z_unconstrained, x_unconstrained)
+ if z_mask is False:
+     log_density = z_unconstrained.new_zeros(())
+ else:
+     log_density = self.transform.log_abs_det_jacobian(z_unconstrained, x_unconstrained)
  self.x_unconstrained = list(reversed(list(self.guide._unpack_latent(x_unconstrained))))