From 2f65617be080694b20e08e2c9062ccbe37c542ac Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Fri, 15 May 2020 13:57:16 +0200 Subject: [PATCH] fix name clash --- pts/modules/distribution_output.py | 24 ++++++++++-------------- 1 file changed, 10 insertions(+), 14 deletions(-) diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index 3504068..2cbe753 100644 --- a/pts/modules/distribution_output.py +++ b/pts/modules/distribution_output.py @@ -191,16 +191,15 @@ class StudentTMixtureOutput(DistributionOutput): def distribution( self, distr_args, scale: Optional[torch.Tensor] = None ) -> Distribution: - mix_logits, df, loc, scale = distr_args + mix_logits, df, loc, dist_scale = distr_args - comp_distr = StudentT(df, loc, scale) + distr = MixtureSameFamily( + Categorical(logits=mix_logits), StudentT(df, loc, dist_scale) + ) if scale is None: - return MixtureSameFamily(Categorical(logits=mix_logits), comp_distr) + return distr else: - scaled_comp_distr = TransformedDistribution( - comp_distr, [AffineTransform(loc=0, scale=scale)] - ) - return MixtureSameFamily(Categorical(logits=mix_logits), scaled_comp_distr) + return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) @property def event_shape(self) -> Tuple: @@ -224,16 +223,13 @@ class NormalMixtureOutput(DistributionOutput): def distribution( self, distr_args, scale: Optional[torch.Tensor] = None ) -> Distribution: - mix_logits, loc, scale = distr_args + mix_logits, loc, dist_scale = distr_args - comp_distr = Normal(loc, scale) + distr = MixtureSameFamily(Categorical(logits=mix_logits), Normal(loc, dist_scale)) if scale is None: - return MixtureSameFamily(Categorical(logits=mix_logits), comp_distr) + return distr else: - scaled_comp_distr = TransformedDistribution( - comp_distr, [AffineTransform(loc=0, scale=scale)] - ) - return MixtureSameFamily(Categorical(logits=mix_logits), scaled_comp_distr) + return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) @property def event_shape(self) -> Tuple: