Add option to provide your own PDF for the latent variables in FlowSequential (not necessarily Gaussian)

[João M. Pereira]

In our application, we've found the need to provide the PDF for the latent variables (in our case, it is useful that the PDF depends on the conditional inputs). We've found a small fix to your code that let's you have this option, by just adding an init method to FlowSequential and changing log_probs accordingly. I'll share the different code here. Here is the added init method

    def __init__(self, *args, log_probs = 'gaussian'):
        super(FlowSequential,self).__init__(*args)
        if log_probs = 'gaussian'
            def __log_probs(x, *_):
                return torch.sum(-0.5 * x ** 2 - 0.5 * math.log(2 * math.pi),
                                                            -1, keepdim=True)
            self.log_probs = __log_probs
        else:
            self.log_probs = log_probs

and here is the modified log_probs.

    def log_probs(self, inputs, cond_inputs = None):
        u, log_jacob = self(inputs, cond_inputs)
        log_probs = self.log_probs(u, cond_inputs)
        return (log_probs + log_jacob).sum(-1, keepdim=True)

Then, to provide your own PDF, just do
model = FlowSequential(*modules,logprobs=your_own_PDF)