Affine Transformation Module

Is there a built-in affine transformation module? Something of the form y = A @ x + b? I see that there is a lower Cholesky affine transformation, but it is not a module. Tensorflow has the sort of thing I’m looking for (their Affine Bijector), and I was a little surprised it doesn’t seem to exist in Pyro.

Hi @kernelmethod

I’m pretty sure there is an active Pull Request in pryo for this now. Does this fit your goal?

Hi, most common transformations live upstream in PyTorch Distributions (torch.distributions.transforms). There’s an AffineTransform there: Probability distributions - torch.distributions — PyTorch 1.13 documentation

Thanks for the replies @JamesTrick and @eb8680_2. Unfortunately, the pull-request wasn’t quite what I was look for, and it was actually rejected shortly after you posted, but I appreciate the pointer!

As for the AffineTransform in the PyTorch Distributions package, the main issue is that PyTorch’s Transforms aren’t trainable; they don’t subclass nn.Module. Pyro has a workaround for this, which provides a class that deals with the persnickety __hash__ function issue. I’d happily just throw together an implementation of AffineTransform subclassing Pyro’s TransformModule, but I was trying to avoid it because inverting affine transformations is an expensive operation unless you use some tricks that are a bit tougher to implement correctly.

Anyway, I was hoping there was an existing implementation of an affine transform that subclasses Pyro’s TransformModule. If there isn’t, maybe I should implement one and create a pull request myself?

Thanks again for the responses!