wassname/pytorch-ts
1
0
Fork 0
mirror of https://github.com/wassname/pytorch-ts.git synced 2026-08-11 11:24:31 +08:00
Code Issues Packages Projects Releases Wiki Activity

fix scaling of mixtures

Browse Source
This commit is contained in:
Dr. Kashif Rasul
2020-05-15 13:21:52 +02:00
parent af0b129b21
commit cd5b061b9b
1 changed files with 12 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pts/modules/distribution_output.py
+12 -8
View File
@@ -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: