Fix multinomial_goodness_of_fit printing (#3115)

pyro/distributions/testing/gof.py

@@ -70,10 +70,10 @@ class InvalidTest(ValueError):
 
 
 def print_histogram(probs, counts):
-    max_count = max(counts)
+    max_count = int(max(counts))
     print("{: >8} {: >8}".format("Prob", "Count"))
     for prob, count in sorted(zip(probs, counts), reverse=True):
-        width = int(round(HISTOGRAM_WIDTH * count / max_count))
+        width = int(round(HISTOGRAM_WIDTH * int(count) / max_count))
         print("{: >8.3f} {: >8d} {}".format(prob, count, "-" * width))
 
 

tests/distributions/testing/test_gof.py

@@ -0,0 +1,20 @@
+# Copyright Contributors to the Pyro project.
+# SPDX-License-Identifier: Apache-2.0
+
+import torch
+import torch.distributions as dist
+
+from pyro.distributions.testing.gof import multinomial_goodness_of_fit
+
+
+def test_multinomial_goodness_of_fit():
+    N = 100000
+    K = 20
+    logits = torch.randn(K)
+    probs = (logits - logits.logsumexp(-1)).exp()
+    d = dist.Categorical(probs)
+    samples = d.sample((N,))
+    counts = torch.zeros(K, dtype=torch.long)
+    counts.scatter_add_(0, samples, torch.ones(N, dtype=torch.long))
+    gof = multinomial_goodness_of_fit(probs, counts, plot=True)
+    assert gof > 0.1