Repository: pyro-ppl/numpyro

Issue:
Bug in Kumaraswamy distribution?
Below is a fully self-contained snippet that specifies a model and (manual) guide and then does forward-simulation and MLE+SVI. The snippet runs through and seems to do what it should be doing. 

However, when I simply replace the beta sample distribution in the guide with a Kumaraswamy distribution (both support in (0, 1), both with two parameters > 0), I receive an error message: "TypeError: log_prob() takes 2 positional arguments but 3 were given". 

This seems like a bug to me?

```python
import numpy as np
import pandas as pd
import jax
import numpyro

def model(disease_a=None, risk_pred_data=None):
    
    disease_a = numpyro.param("disease_a", lambda rng_key: jax.random.gamma(rng_key, 2) if not disease_a else disease_a, 
                              constraint=numpyro.distributions.constraints.positive)

    if risk_pred_data is not None:
        nobs = len(risk_pred_data)            
    else:
        nobs = 1

    with numpyro.plate("N", nobs):
        # latent state in (0, 1)
        # basic reparametrization of Kumaraswamy(a,1) in terms of Unif(0,1).
        disease_base = numpyro.sample("disease_base", numpyro.distributions.Uniform(np.zeros((nobs,)) + 1e-3, 
                                                                                    np.ones((nobs,)) - 1e-3))
        disease_state = numpyro.deterministic("disease_state", jax.lax.clamp(1e-3, disease_base**(1/disease_a), 1-1e-3))

        risk_pred = numpyro.sample("risk_pred", numpyro.distributions.BetaProportion(mean=disease_state, concentration=1),
                                   obs=None if risk_pred_data is None else jax.lax.clamp(1e-3, risk_pred_data, 1-1e-3))

    
def guide(disease_a=None, risk_pred_data=None):
    
    disease_a = numpyro.param("disease_a", lambda rng_key: jax.random.gamma(rng_key, 2) if not disease_a else disease_a,
                              constraint=numpyro.distributions.constraints.positive)

    if risk_pred_data is not None:
        nobs = len(risk_pred_data)            
    else:
        nobs = 1

    disease_base_a = numpyro.param("disease_base_a", lambda rng_key: jax.random.gamma(rng_key, 2*np.ones((nobs,))),
                                   constraint=numpyro.distributions.constraints.positive)
    disease_base_b = numpyro.param("disease_base_b", lambda rng_key: jax.random.gamma(rng_key, 2*np.ones((nobs,))),
                                   constraint=numpyro.distributions.constraints.positive)

    with numpyro.plate("N", nobs):    
        # works
        disease_base = numpyro.sample("disease_base", numpyro.distributions.Beta(disease_base_a, disease_base_b))
        # raises an error
        # disease_base = numpyro.sample("disease_base", numpyro.distributions.Kumaraswamy(disease_base_a, disease_base_b))

# forward-sim
mdl = lambda: model(disease_a=1.0)
data_sim = numpyro.infer.Predictive(mdl, num_samples=10000)(jax.random.PRNGKey(2))
df_sim = pd.DataFrame({k: v.reshape(-1) for (k, v) in data_sim.items()})

# MLE/SVI
optimizer = numpyro.optim.Adam(step_size=0.0001)
svi = numpyro.infer.SVI(model, guide, optimizer, loss=numpyro.infer.autoguide.Trace_ELBO())
svi_result, last_state, losses = svi.run(jax.random.PRNGKey(0), 1000, risk_pred_data=df_sim.risk_pred.to_numpy().squeeze())
```

The error occurs in 
https://github.com/pyro-ppl/numpyro/blob/bcf38d7e2ca97fded876cf034866eaa0f00385ff/numpyro/distributions/util.py#L675 called from https://github.com/pyro-ppl/numpyro/blob/bcf38d7e2ca97fded876cf034866eaa0f00385ff/numpyro/infer/util.py#L71 which is called in https://github.com/pyro-ppl/numpyro/blob/bcf38d7e2ca97fded876cf034866eaa0f00385ff/numpyro/infer/elbo.py#L129

I am using numpyro 0.13.2.


Discussion:
Thanks, this is indeed a bug. Could you try adding `intermediates=None` to the signature of Kumaraswamy.log_prob https://github.com/pyro-ppl/numpyro/blob/aeeeb72994d7e756abdf40b1b47a7b4bc929f4bf/numpyro/distributions/continuous.py#L806
```
def log_prob(self, value, intermediates=None):
```
Thanks for the quick response! That indeed fixes the TypeError issue. 

There seems to be some other problem remaining with the Kumaraswamy distribution, however. (Probably completely unrelated from the above issue.)

For a=1, [we have](https://en.wikipedia.org/wiki/Kumaraswamy_distribution) that Kumaraswamy(1, b) = Beta(1, b). So when I fix the `disease_base_a` params in the guide to be all ones (`np.ones((nobs,))`, no random initialization), switching between Beta and Kumaraswamy should not make any difference at all, I believe?

What I observe, however, is that SVI runs smoothly with the Beta distribution but returns initial loss inf and subsequent losses nan for the Kumaraswamy distribution. 

I did some preliminary investigations, and it seems to me that the posterior samples and log likelihood look fine. So my best guess would be that the issue stems from the KL implementation for the Kumaraswamy? I tried setting KL_KUMARASWAMY_BETA_TAYLOR_ORDER = 1000, which did not seem to change anything. 
I'm not sure why I made Kumaraswamy a TransformedDistribution. The issue is TransformedDistribution.sample method does not deal with boundary values of the support properly. You can simply make Kumaraswamy a normal distribution like
```python
class Kumaraswamy(Distribution):
        ...
        batch_shape = lax.broadcast_shapes(
            jnp.shape(concentration1), jnp.shape(concentration0)
        )
        super().__init__(batch_shape=batch_shape, validate_args=validate_args)
```
Could you submit a PR for this?
Yes, I'll do that! 

I still get somewhat different numerical behavior (e.g. initial loss 27742.4551 with Beta vs 28135.0059 with Kuma), but I suspect that is expected?
yes, it is expected. samples generated from two distributions might come from different random paths (up to the implementation).
I believe there is still something else going on with the Kumaraswamy distribution... I have another case where if I use the Kuma distribution, I get a lot of jax nan errors while none of those occur when using the beta distribution instead. 

They are of course different distributions and thus doing different things during inference, but the fact that Kuma raises nan errors while beta does not (in an otherwise identical setup, especially as it relates to their parameters) seems to indicate another bug to me. 

I tried to track down why exactly the error occurs, but was unable to nail it down. The error ("FloatingPointError: invalid value (nan) encountered in jit(mul)") occurs in [this line](https://github.com/pyro-ppl/numpyro/blob/1528036d0023a37fe3fd9934365f48f11be3caab/numpyro/distributions/continuous.py#L797). You can find the reproducible setup [here ](https://gist.github.com/e-pet/ffe6594fc94ecf191519060009e03bf9)- the Kuma version produces the first nan errors in iteration 4 for me; the beta version doesn't (at all, as far as I can see). Sorry for the somewhat complex model setup - I tried to simplify further, but then the nan errors would always go away. I'm sure there's a simpler setup that produces the error; I just haven't found it yet.

Resolving patch:
diff --git a/numpyro/distributions/continuous.py b/numpyro/distributions/continuous.py
--- a/numpyro/distributions/continuous.py
+++ b/numpyro/distributions/continuous.py
@@ -790,18 +790,22 @@ def __init__(self, concentration1, concentration0, *, validate_args=None):
 
     def sample(self, key, sample_shape=()):
         assert is_prng_key(key)
-        minval = jnp.finfo(jnp.result_type(float)).tiny
-        u = random.uniform(key, shape=sample_shape + self.batch_shape, minval=minval)
-        log_sample = jnp.log1p(-(u ** (1 / self.concentration0))) / self.concentration1
-        finfo = jnp.finfo(u)
+        finfo = jnp.finfo(jnp.result_type(float))
+        u = random.uniform(
+            key, shape=sample_shape + self.batch_shape, minval=finfo.tiny
+        )
+        u_con0 = jnp.clip(u ** (1 / self.concentration0), a_max=1 - finfo.eps)
+        log_sample = jnp.log1p(-u_con0) / self.concentration1
         return jnp.clip(jnp.exp(log_sample), a_min=finfo.tiny, a_max=1 - finfo.eps)
 
     @validate_sample
     def log_prob(self, value):
+        finfo = jnp.finfo(jnp.result_type(float))
         normalize_term = jnp.log(self.concentration0) + jnp.log(self.concentration1)
+        value_con1 = jnp.clip(value**self.concentration1, a_max=1 - finfo.eps)
         return (
             xlogy(self.concentration1 - 1, value)
-            + xlog1py(self.concentration0 - 1, -(value**self.concentration1))
+            + xlog1py(self.concentration0 - 1, -value_con1)
             + normalize_term
         )
 


Test patch: