PyTorch function transforms are not compatible with PyroParams

[eb8680]

In chirho.robust, we're making extensive use of torch.func, especially the vectorization transform torch.func.vmap and forward and reverse-mode autodiff transforms torch.func.jvp/vjp.

Unfortunately, the autodiff transforms seem to be fundamentally incompatible with PyroModules that have constrained PyroParams, including most of the standard autoguides in pyro.infer.autoguide - the raw PyTorch Parameter tensors underlying the constrained PyroParams do not seem to receive or propagate gradients. It's not clear what exactly in Pyro is causing the problem or how to fix or work around it, although it at least doesn't seem specific to ChiRho.

This isn't immediately blocking other work, but without compatibility with PyroModules the practical utility of chirho.robust will be significantly diminished, so we should try to resolve this somehow before merging #398 into master.

[agrawalraj]

Here are some odd interactions of torch.func:

import torch


def f1(x):
    return x**2


def f2(x):
    return torch.tensor([x[0] ** 2, x[1] ** 2, x[2] ** 2])


x = 1.0 * torch.ones(3, requires_grad=True)
v = torch.ones(3)

print(torch.func.jvp(f1, (x,), (v,))[1]) # returns tensor([2., 2., 2.], grad_fn=<MulBackward0>)
print(torch.func.jvp(f2, (x,), (v,))[1]) # returns tensor([0., 0., 0.])

Figured I would post in case it was related at all to the composability issue with PyroParams.

[eb8680]

This should be resolved once a version of Pyro including pyro-ppl/pyro#3328 is released.

[agrawalraj]

Ok that's awesome, thanks!

[eb8680]

Resolved in the Pyro 1.9.0 release