From cd5b061b9b09e2b4d17ad7599f241a79ae8c0b50 Mon Sep 17 00:00:00 2001 From: "Dr. Kashif Rasul" Date: Fri, 15 May 2020 13:21:52 +0200 Subject: [PATCH] fix scaling of mixtures --- pts/modules/distribution_output.py | 20 ++++++++++++-------- 1 file changed, 12 insertions(+), 8 deletions(-) diff --git a/pts/modules/distribution_output.py b/pts/modules/distribution_output.py index ae2166d..3504068 100644 --- a/pts/modules/distribution_output.py +++ b/pts/modules/distribution_output.py @@ -193,13 +193,14 @@ class StudentTMixtureOutput(DistributionOutput): ) -> Distribution: mix_logits, df, loc, scale = distr_args - distr = MixtureSameFamily( - Categorical(logits=mix_logits), StudentT(df, loc, scale) - ) + comp_distr = StudentT(df, loc, scale) if scale is None: - return distr + return MixtureSameFamily(Categorical(logits=mix_logits), comp_distr) else: - return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) + scaled_comp_distr = TransformedDistribution( + comp_distr, [AffineTransform(loc=0, scale=scale)] + ) + return MixtureSameFamily(Categorical(logits=mix_logits), scaled_comp_distr) @property def event_shape(self) -> Tuple: @@ -225,11 +226,14 @@ class NormalMixtureOutput(DistributionOutput): ) -> Distribution: mix_logits, loc, scale = distr_args - distr = MixtureSameFamily(Categorical(logits=mix_logits), Normal(loc, scale)) + comp_distr = Normal(loc, scale) if scale is None: - return distr + return MixtureSameFamily(Categorical(logits=mix_logits), comp_distr) else: - return TransformedDistribution(distr, [AffineTransform(loc=0, scale=scale)]) + scaled_comp_distr = TransformedDistribution( + comp_distr, [AffineTransform(loc=0, scale=scale)] + ) + return MixtureSameFamily(Categorical(logits=mix_logits), scaled_comp_distr) @property def event_shape(self) -> Tuple: