Repository: pyro-ppl/numpyro

Issue:
Sampling from posterior predictive with Categorical and AutoDiagonalNormal occasionally samples outside support
We encountered a subtle problem when sampling from the posterior predictive distribution for a model containing `Categorical` and using the `AutoDiagonalNormal`: In some rare cases, the sample drawn from the `Categorical` would be one large than the maximum category index. E.g., a model like
```
def model(num_categories):
    sample("ps", Dirichlet(jnp.ones(num_categories)))
    sample("x", Categorical(probs=ps))
```
may sample a value of `num_categories` even though the support of the categorical distribution is `(0, ..., num_categories-1)`.

This seems to arise from the combination of
- `numpyro.distributions.util._categorical` (and hence `numpyro.distributions.CategoricalProbs`) [expecting the passed probability vector to sum up exactly to one](https://github.com/pyro-ppl/numpyro/blob/1dbe63f72f1cd8befb76d0b8e7b4f284e7852404/numpyro/distributions/util.py#L187)
- `AutoDiagonalNormal` (or probably rather the `StickBreakingTransform`) not guaranteeing that due to numerical imprecision

We observed samples from the `AutoDiagonalNormal` guide for which the probability vector `ps` summed up to something just shy of one. In rare cases, the uniform (0,1) sample used in `numpyro.distributions.util._categorical` falls above the total sum, resulting in the function to return the largest category index plus one as a value.

Our current quick-and-dirty workaround for this is to simply clamp the sampled values to the support but surely there are better ways of handling this, but I haven't given it much thought. Naively it would seem that explicit renormalisation of the probability vector at some point would be the way to go, but I'm not sure that will actually be robust towards the numerical issues that seem to lead to this in the first place..


Discussion:
> the uniform (0,1) sample used in numpyro.distributions.util._categorical falls above the total sum, resulting in the function to return the largest category index plus one as a value.

Good point! I think one way is to remove the last value `1` there: `s = s[..., :-1]`. Do you want to submit the fix? :)
I'll give that a go soon.

Edit: Thinking about it a bit more, the solution you suggest would introduce as slight bias towards the last category, which would completely subsume the error/difference of probabilities summing up to one. Given that, I would prefer renormalising s, i.e.
`s = s / s[..., -1]` which would keep the relative proportions of the probabilities intact - that seems to be the solution with less surprise factor for me (and contrary to my earlier rambling about potential issues would be perfectly fine to mitigate the issue). Thoughts?
This is also the reason for tests failing in funsor https://github.com/pyro-ppl/funsor/issues/594. Funsor assumes that Categorical distribution `probs` are renormalized by the backend distribution which is done by the Pytorch's Categorical distribution but apparently not in Numpyro.
> prefer renormalising s, i.e. s = s / s[..., -1]

Yeah, please go with it. Maybe we can also do `s = s[..., :-1] / s[..., -1]` if the division is not perfect (I don't know)
@ordabayevy Thanks for the pointer! We typically don't want to modify the inputs of the distributions in numpyro so I hope that @lumip solution will resolve the issue. Maybe we can also fix the issue [here](https://github.com/pyro-ppl/numpyro/blob/6e98655957196243c2ab178e3af861060172cf4d/numpyro/distributions/discrete.py#L68) for completeness:
```
def _to_logits_multinom(probs):
    minval = jnp.finfo(jnp.result_type(probs)).min
    return jnp.clip(jnp.log(probs) - logsumexp(probs, axis=-1, keepdims=True), a_min=minval)
```

Resolving patch:
diff --git a/numpyro/distributions/util.py b/numpyro/distributions/util.py
--- a/numpyro/distributions/util.py
+++ b/numpyro/distributions/util.py
@@ -185,6 +185,8 @@ def _categorical(key, p, shape):
     # Ref: https://stackoverflow.com/a/34190035
     shape = shape or p.shape[:-1]
     s = jnp.cumsum(p, axis=-1)
+    # Normalize s to deal with numerical issues.
+    s = s[..., :-1] / s[..., -1:]
     r = random.uniform(key, shape=shape + (1,))
     # FIXME: replace this computation by using binary search as suggested in the above
     # reference. A while_loop + vmap for a reshaped 2D array would be enough.
diff --git a/setup.py b/setup.py
--- a/setup.py
+++ b/setup.py
@@ -63,6 +63,7 @@
             "graphviz",
             "jaxns==1.0.0",
             "optax>=0.0.6",
+            "pyyaml",  # flax dependency
             "tensorflow_probability>=0.15.0",
         ],
         "examples": [


Test patch: