From e23e363bf58f821e1266bf2d6c1a9cea3c4ec377 Mon Sep 17 00:00:00 2001 From: Bryce Date: Wed, 21 Sep 2022 09:14:39 -0700 Subject: [PATCH] lint --- imaginairy/modules/distributions.py | 33 +++++++++++++++-------------- 1 file changed, 17 insertions(+), 16 deletions(-) diff --git a/imaginairy/modules/distributions.py b/imaginairy/modules/distributions.py index 1cb8694..62a671b 100644 --- a/imaginairy/modules/distributions.py +++ b/imaginairy/modules/distributions.py @@ -24,23 +24,24 @@ class DiagonalGaussianDistribution: def kl(self, other=None): if self.deterministic: return torch.Tensor([0.0]) - else: - if other is None: - return 0.5 * torch.sum( - torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, - dim=[1, 2, 3], - ) - else: - return 0.5 * torch.sum( - torch.pow(self.mean - other.mean, 2) / other.var - + self.var / other.var - - 1.0 - - self.logvar - + other.logvar, - dim=[1, 2, 3], - ) - def nll(self, sample, dims=[1, 2, 3]): + if other is None: + return 0.5 * torch.sum( + torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, + dim=[1, 2, 3], + ) + + return 0.5 * torch.sum( + torch.pow(self.mean - other.mean, 2) / other.var + + self.var / other.var + - 1.0 + - self.logvar + + other.logvar, + dim=[1, 2, 3], + ) + + def nll(self, sample, dims=None): + dims = dims if dims is None else [1, 2, 3] if self.deterministic: return torch.Tensor([0.0]) logtwopi = np.log(2.0 * np.pi)