Normalizing flows, transforms as Funsors?

[fritzo]

How would we implement normalizing flows / Transforms in Funsor?

As a first guess, each transform would be a funsor with one free variable of the same domain as its output, and we would have a .inv attribute like torch.distributions.Transform.inv to get the inverse transform. Additionally there would be rules for

• eager substitution t(x) which simply calls the forward transform
• inverse followed by eager substitution t.inv(x) which calls the inverse transform
• eager substitution of Deltas

  Delta("x", x, log_density)(x=t("y"))
      => Delta("y", t.inv(x), log_density - t.log_abs_det_jacobian(x))

• eager reduce of a substitution

  f(x=t(y)).reduce(ops.logaddexp, "y")
      => (f(x) + t.inv.log_abs_det_jacobian(x)).reduce(ops.logaddexp, "x")

• other reduction rules?

I think we could implement this as a single class that took a PyTorch/NumPyro transform object

class Transform(Funsor):  # or maybe Bijector?
    "Funsor backed by a PyTorch transform object."
    def __init__(self, transform):
        self.transform = transform
    @property
    def inv(self):
        return Transform(self.transform.inv)
    def eager_subs(self, subs):
        ...
        return Tensor(self.transform(x.data))
    ...

cc @stefanwebb

[fritzo]

Note @eb8680 and I sketched one possible direction in https://github.com/pyro-ppl/funsor/tree/lebesgue

[eb8680]

I'm going to close this as resolved by #365, #427 and #442. We can open a new issue if we decide to revisit the design.