bound negbin scale above 1 (awslabs#814)

kashif · May 10, 2020 · 46e22a7 · 46e22a7
1 parent b42a85f
commit 46e22a7
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/src/gluonts/distribution/neg_binomial.py b/src/gluonts/distribution/neg_binomial.py
@@ -129,6 +129,7 @@ def distribution(
             return NegativeBinomial(mu, alpha)
         else:
             F = getF(mu)
+            scale = 1.0 + softplus(F, scale - 1.0)
             mu = F.broadcast_mul(mu, scale)
             alpha = F.broadcast_add(alpha, F.broadcast_div(scale - 1, mu))
             return NegativeBinomial(mu, alpha, F)