Repository: pyro-ppl/numpyro

Issue:
Vectorized sampling: error when var size is 0
JAX refuses to reshape an array using the `(-1, ...)` shape notation if its `size` (i.e. product of `shape`) is 0. This causes issues with vectorized sampling when samples from different chains are concatenated.

# Steps to reproduce
```python
import numpy as np
import numpyro

import numpyro.distributions as dist

from jax import random
from numpyro.infer import MCMC, NUTS

J = 8
n = 0
y = np.random.rand(n, J)
sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])


# Slightly modified Eight Schools example

def eight_schools(J, sigma, y=None):
    mu = numpyro.sample('mu', dist.Normal(0, 5))
    tau = numpyro.sample('tau', dist.HalfCauchy(5))
    with numpyro.plate('J', J):
        theta = numpyro.sample('theta', dist.Normal(mu, tau))
        numpyro.sample('obs', dist.Normal(theta, sigma), sample_shape=(n,))

nuts_kernel = NUTS(eight_schools)

mcmc = MCMC(nuts_kernel, num_warmup=50, num_samples=100, num_chains=2, chain_method='vectorized')

rng_key = random.PRNGKey(0)

mcmc.run(rng_key, J, sigma, y=y, extra_fields=('potential_energy',))
```

Sampling of individual chains works fine but a `[TypeError: reshape new_sizes must all be positive, got (-1, 0, 8).]()` is raised during final concatenation of samples.

# Possible fix
In `mcmc.py:612`

```python
states_flat = tree_map(
    lambda x: jnp.reshape(x, (-1,) + x.shape[2:]), states
)
```

simply calculating the dimension of the first axis is sufficient

```python
states_flat = tree_map(
    lambda x: jnp.reshape(x, (x.shape[0] * x.shape[1],) + x.shape[2:]), states
)
```

It seems unlikely that the issue can be "fixed" upstream in jax, since numpy has the same policy towards reshaping. Happy to open a PR with the above fix if it seems reasonable.


Discussion:
Oh, what's the subtle issue. Please feel free to make the change.

Resolving patch:
diff --git a/numpyro/infer/mcmc.py b/numpyro/infer/mcmc.py
--- a/numpyro/infer/mcmc.py
+++ b/numpyro/infer/mcmc.py
@@ -609,7 +609,9 @@ def run(self, rng_key, *args, extra_fields=(), init_params=None, **kwargs):
                 # swap num_samples x num_chains to num_chains x num_samples
                 states = tree_map(lambda x: jnp.swapaxes(x, 0, 1), states)
             states_flat = tree_map(
-                lambda x: jnp.reshape(x, (-1,) + x.shape[2:]), states
+                # need to calculate first dimension manually; see issue #1328
+                lambda x: jnp.reshape(x, (x.shape[0] * x.shape[1],) + x.shape[2:]),
+                states,
             )
         self._last_state = last_state
         self._states = states


Test patch:
diff --git a/test/infer/test_mcmc.py b/test/infer/test_mcmc.py
--- a/test/infer/test_mcmc.py
+++ b/test/infer/test_mcmc.py
@@ -1080,3 +1080,29 @@ def model():
     mcmc = MCMC(NUTS(model), num_warmup=10, num_samples=10)
     with pytest.warns(FutureWarning, match="enumerated sites"):
         mcmc.run(random.PRNGKey(0))
+
+
+def test_vectorized_sampling_zero_sized():
+    J = 8
+    n = 0
+    sigma = np.array([15.0, 10.0, 16.0, 11.0, 9.0, 11.0, 10.0, 18.0])
+
+    # Modified Eight Schools example
+    def eight_schools(J, sigma):
+        mu = numpyro.sample("mu", dist.Normal(0, 5))
+        tau = numpyro.sample("tau", dist.HalfCauchy(5))
+        with numpyro.plate("J", J):
+            theta = numpyro.sample("theta", dist.Normal(mu, tau))
+            numpyro.sample("scores", dist.Normal(theta, sigma), sample_shape=(n,))
+
+    nuts_kernel = NUTS(eight_schools)
+    mcmc = MCMC(
+        nuts_kernel,
+        num_warmup=50,
+        num_samples=100,
+        num_chains=2,
+        chain_method="vectorized",
+    )
+    rng_key = random.PRNGKey(0)
+    mcmc.run(rng_key, J, sigma, extra_fields=("potential_energy",))
+    assert mcmc.get_samples()["scores"].shape == (200, 0, 8)