Repository: pyro-ppl/pyro

Issue:
"HMC shape error at Dirichlet sample site"
I'm running into an issue when using Dirichlet latent sites in NUTS. 

The following code:

```python 
def wdk_gp_regressor(Ks_train):
    n_k, n_train, _ = Ks_train.size()
    w = pyro.sample('w', dist.Dirichlet(torch.ones(n_k).to(device))).view((n_k, 1, 1))
    alpha = pyro.sample('alpha', dist.Uniform(torch.tensor([0.0]).to(device), 
                                              torch.tensor([30.0]).to(device)))
    

    sigma = pyro.sample('sigma', dist.Uniform(torch.tensor([1e-2]).to(device), 
                                              torch.tensor([1.0]).to(device)))
    # Covariance
    cov = torch.sum((Ks_train * w), dim=0) * alpha 
    cov += torch.eye(n_train).to(device) * (sigma + 1e-4)
    L = torch.potrf(cov, upper=False)

    # Likelihood
    return pyro.sample('f', dist.MultivariateNormal(torch.zeros(n_train).to(device), 
                                                                           scale_tril=L))
    
def conditioned_gp_regressor(gp_regressor, Ks, y_train):
    return pyro.poutine.condition(gp_regressor, data={'f': y_train})(Ks)

nuts_kernel = NUTS(conditioned_gp_regressor, step_size=1.0, adapt_step_size=True)
lact_posterior = MCMC(nuts_kernel, num_samples=1000, warmup_steps=500).run(wdk_gp_regressor, 
                                                                           torch.stack(Ks_train), y_train)
```

Results in:

```
~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/pyro/ops/integrator.py in single_step_velocity_verlet(z, r, potential_fn, step_size, z_grads)
     51 
     52     for site_name in z_next:
---> 53         r_next[site_name] = r_next[site_name] + 0.5 * step_size * (-grads[site_name])
     54         z_next[site_name] = z_next[site_name] + step_size * r_next[site_name]
     55     grads, potential_energy = _grad(potential_fn, z_next)

RuntimeError: The size of tensor a (25) must match the size of tensor b (24) at non-singleton dimension 0
```

For reference, `Ks_train` has size (25, 590, 590). 


Discussion:

Resolving patch:
diff --git a/pyro/infer/mcmc/hmc.py b/pyro/infer/mcmc/hmc.py
--- a/pyro/infer/mcmc/hmc.py
+++ b/pyro/infer/mcmc/hmc.py
@@ -185,11 +185,13 @@ def setup(self, *args, **kwargs):
         if self._automatic_transform_enabled:
             self.transforms = {}
         for name, node in sorted(trace.iter_stochastic_nodes(), key=lambda x: x[0]):
-            r_loc = torch.zeros_like(node["value"])
-            r_scale = torch.ones_like(node["value"])
-            self._r_dist[name] = dist.Normal(loc=r_loc, scale=r_scale)
+            site_value = node["value"]
             if node["fn"].support is not constraints.real and self._automatic_transform_enabled:
                 self.transforms[name] = biject_to(node["fn"].support).inv
+                site_value = self.transforms[name](node["value"])
+            r_loc = site_value.new_zeros(site_value.shape)
+            r_scale = site_value.new_ones(site_value.shape)
+            self._r_dist[name] = dist.Normal(loc=r_loc, scale=r_scale)
         self._validate_trace(trace)
 
         if self.adapt_step_size:


Test patch:
diff --git a/tests/infer/mcmc/test_hmc.py b/tests/infer/mcmc/test_hmc.py
--- a/tests/infer/mcmc/test_hmc.py
+++ b/tests/infer/mcmc/test_hmc.py
@@ -205,7 +205,6 @@ def model(data):
     assert_equal(posterior.mean, true_std, prec=0.05)
 
 
-@pytest.mark.xfail(reason='log_abs_det_jacobian not implemented for StickBreakingTransform')
 def test_categorical_dirichlet():
     def model(data):
         concentration = torch.tensor([1.0, 1.0, 1.0])
diff --git a/tests/infer/mcmc/test_nuts.py b/tests/infer/mcmc/test_nuts.py
--- a/tests/infer/mcmc/test_nuts.py
+++ b/tests/infer/mcmc/test_nuts.py
@@ -149,3 +149,19 @@ def model(data):
     mcmc_run = MCMC(nuts_kernel, num_samples=500, warmup_steps=100).run(data)
     posterior = EmpiricalMarginal(mcmc_run, sites="p_latent")
     assert_equal(posterior.mean, true_probs, prec=0.03)
+
+
+@pytest.mark.filterwarnings("ignore:Encountered NAN")
+def test_categorical_dirichlet():
+    def model(data):
+        concentration = torch.tensor([1.0, 1.0, 1.0])
+        p_latent = pyro.sample('p_latent', dist.Dirichlet(concentration))
+        pyro.sample("obs", dist.Categorical(p_latent), obs=data)
+        return p_latent
+
+    true_probs = torch.tensor([0.1, 0.6, 0.3])
+    data = dist.Categorical(true_probs).sample(sample_shape=(torch.Size((2000,))))
+    nuts_kernel = NUTS(model, adapt_step_size=True)
+    mcmc_run = MCMC(nuts_kernel, num_samples=200, warmup_steps=100).run(data)
+    posterior = EmpiricalMarginal(mcmc_run, sites='p_latent')
+    assert_equal(posterior.mean, true_probs, prec=0.02)