Repository: pyro-ppl/numpyro

Issue:
initialize_model() fails with batch size of 1 for model with discrete variables
Hi,

I'm having issues using NumPyro's latest version of `initialize_model()` on a model with discrete variables. The model samples random variables within a plate context with a data-dependent size. When the given data is of size $N > 1$, everything works fine. However, if I pass a dataset with a single data point ($N = 1$), the execution of `initialize_model` fails. Apparently, the funsor API expected the computed log-probability of a sampled variable to be a scalar, when in fact it's a single-element 1D array for this case.

Minimal working example:
```python
import jax
import jax.numpy as jnp
import numpyro
from numpyro.distributions import Bernoulli, MultivariateNormal
from numpyro.infer.util import initialize_model


def model(X: jnp.ndarray, y: jnp.ndarray, z: jnp.ndarray):
    phi = numpyro.sample('phi', MultivariateNormal(jnp.zeros(2), jnp.eye(2)))

    n_data = X.shape[-2]
    with numpyro.plate('individual', n_data):
        weights = jnp.tensordot(X, phi, axes=1)
        numpyro.sample('labels', Bernoulli(probs=(1/(1 + jnp.exp(-weights)))), obs=z, infer={'enumerate': 'parallel'})
        coefficients = numpyro.sample('coefficients', MultivariateNormal(jnp.zeros(2), jnp.eye(2)))
        numpyro.sample('responses', MultivariateNormal(X * coefficients, jnp.eye(2)), obs=y)


if __name__ == '__main__':
    rng_key = jax.random.PRNGKey(0)

    # This line runs without problems:
    res = initialize_model(rng_key, model, model_kwargs={'X': jnp.zeros((10, 2)), 'z': None, 'y': None})

    # This line fails
    res1 = initialize_model(rng_key, model, model_kwargs={'X': jnp.zeros((1, 2)), 'z': None, 'y': None})

    print('Done')
```

Execution output:
```
Traceback (most recent call last):
  File "/home/rafael/Projects/trajectory-clustering/funsors_issue.py", line 26, in <module>
    res1 = initialize_model(rng_key, model, model_kwargs={'X': jnp.zeros((1, 2)), 'z': None, 'y': None})
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 654, in initialize_model
    (init_params, pe, grad), is_valid = find_valid_initial_params(
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 395, in find_valid_initial_params
    (init_params, pe, z_grad), is_valid = _find_valid_params(
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 381, in _find_valid_params
    _, _, (init_params, pe, z_grad), is_valid = init_state = body_fn(init_state)
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 366, in body_fn
    pe, z_grad = value_and_grad(potential_fn)(params)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/_src/api.py", line 1063, in value_and_grad_f
    ans, vjp_py = _vjp(f_partial, *dyn_args, reduce_axes=reduce_axes)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/_src/api.py", line 2558, in _vjp
    out_primal, out_vjp = ad.vjp(
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/interpreters/ad.py", line 133, in vjp
    out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/interpreters/ad.py", line 122, in linearize
    jaxpr, out_pvals, consts = pe.trace_to_jaxpr_nounits(jvpfun_flat, in_pvals)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/_src/profiler.py", line 312, in wrapper
    return func(*args, **kwargs)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 621, in trace_to_jaxpr_nounits
    jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/jax/linear_util.py", line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 248, in potential_energy
    log_joint, model_trace = log_density_(
  File "/home/rafael/Projects/numpyro/numpyro/contrib/funsor/infer_util.py", line 270, in log_density
    result, model_trace, _ = _enum_log_density(
  File "/home/rafael/Projects/numpyro/numpyro/contrib/funsor/infer_util.py", line 181, in _enum_log_density
    log_prob_factor = funsor.to_funsor(
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/functools.py", line 888, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/funsor/tensor.py", line 491, in tensor_to_funsor
    raise ValueError(
jax._src.traceback_util.UnfilteredStackTrace: ValueError: Invalid shape: expected (), actual (1,)

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/home/rafael/Projects/numpyro/numpyro/infer/util.py", line 248, in potential_energy
    log_joint, model_trace = log_density_(
  File "/home/rafael/Projects/numpyro/numpyro/contrib/funsor/infer_util.py", line 270, in log_density
    result, model_trace, _ = _enum_log_density(
  File "/home/rafael/Projects/numpyro/numpyro/contrib/funsor/infer_util.py", line 181, in _enum_log_density
    log_prob_factor = funsor.to_funsor(
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/functools.py", line 888, in wrapper
    return dispatch(args[0].__class__)(*args, **kw)
  File "/home/rafael/opt/anaconda3/envs/numpyro_dev/lib/python3.9/site-packages/funsor/tensor.py", line 491, in tensor_to_funsor
    raise ValueError(
ValueError: Invalid shape: expected (), actual (1,)
```


Discussion:
Thanks @rafaol! I guess we mixed up 1-size plate with singleton dimensions created by promoting and broadcasting somewhere. Let us take a closer look into this.
@fehiepsi I think I uncovered the source of the bug. 

<details>

In this line of `_enum_log_density`, we try to cast the log_prob_factor to a funsor, defaulting to a funsor of shape `()` if there is no `dim_to_name` argument

https://github.com/pyro-ppl/numpyro/blob/7c3ec50eff96f4bc56fa945fd73f1dee32395888/numpyro/contrib/funsor/infer_util.py#L222-L225 

Over in `NamedMessenger`, we don't make a name request if the batch shape is 1 and if there isn't an existing name, which is exactly the case for plates of size 1.  

https://github.com/pyro-ppl/numpyro/blob/7c3ec50eff96f4bc56fa945fd73f1dee32395888/numpyro/contrib/funsor/enum_messenger.py#L275-L285

As a result `enum_log_density` tries to cast the log_prob associated with the variable (of shape `(1,)`) in the plate of size 1  into the Real funsor with shape `()` and fails.

Here's a minimal example showing the difference in dim_to_name allocation depending on plate size:
```python
import numpyro
import numpyro.distributions as dist
from numpyro.contrib.funsor.infer_util import plate_to_enum_plate
from numpyro.contrib.funsor.enum_messenger import enum
from numpyro.contrib.funsor.enum_messenger import (
    infer_config,
    plate as enum_plate,
    trace as packed_trace,
)

def working_model():
    with numpyro.plate('n', 2):
        numpyro.sample('a', dist.Normal())

def failing_model():
    with numpyro.plate('n', 1):
        numpyro.sample('a', dist.Normal())

with enum(first_available_dim=-2):
    with plate_to_enum_plate():
        good_trace = packed_trace(numpyro.handlers.seed(working_model, rng_seed=0)).get_trace()
        
with enum(first_available_dim=-2):
    with plate_to_enum_plate():
        bad_trace = packed_trace(numpyro.handlers.seed(failing_model, rng_seed=0)).get_trace()

print(good_trace['a']['infer']['dim_to_name'], good_trace['a']['infer']['name_to_dim'])
print(bad_trace['a']['infer']['dim_to_name'], bad_trace['a']['infer']['name_to_dim'])
``` 
</details>
A quick fix is to insert this code right before trying to cast to funsor. 

```python
if log_prob.shape == (1,) and dim_to_name == OrderedDict():
   log_prob = log_prob.squeeze()
```


Thanks, @amifalk! Do you know what happens if log_prob shape has non-trivial singleton dimension, e.g., (2, 1, 3), (2, 1), (1, 2), (1, 1)?
It only fails when every single plate dimension for the variable is 1. i.e.  (2, 1, 3), (2, 1), (1, 2) work, but (1, 1, ..., 1) fails. So I guess we just need to check that `all(dim == 1 for dim in log_prob.shape)` 
Woohoo, that's a good news. Thanks for digging into this @amifalk! Could you send a PR?

Resolving patch:
diff --git a/numpyro/contrib/funsor/infer_util.py b/numpyro/contrib/funsor/infer_util.py
--- a/numpyro/contrib/funsor/infer_util.py
+++ b/numpyro/contrib/funsor/infer_util.py
@@ -1,7 +1,7 @@
 # Copyright Contributors to the Pyro project.
 # SPDX-License-Identifier: Apache-2.0
 
-from collections import defaultdict
+from collections import OrderedDict, defaultdict
 from contextlib import contextmanager
 import functools
 import re
@@ -220,6 +220,10 @@ def _enum_log_density(model, model_args, model_kwargs, params, sum_op, prod_op):
                 log_prob = scale * log_prob
 
             dim_to_name = site["infer"]["dim_to_name"]
+
+            if all(dim == 1 for dim in log_prob.shape) and dim_to_name == OrderedDict():
+                log_prob = log_prob.squeeze()
+
             log_prob_factor = funsor.to_funsor(
                 log_prob, output=funsor.Real, dim_to_name=dim_to_name
             )


Test patch:
diff --git a/test/contrib/test_funsor.py b/test/contrib/test_funsor.py
--- a/test/contrib/test_funsor.py
+++ b/test/contrib/test_funsor.py
@@ -575,3 +575,18 @@ def gmm(data):
     with pytest.raises(Exception):
         mcmc.run(random.PRNGKey(2), data)
     assert len(_PYRO_STACK) == 0
+
+
+@pytest.mark.parametrize(
+    "i_size, j_size, k_size", [(1, 1, 1), (1, 2, 1), (2, 1, 1), (1, 1, 2)]
+)
+def test_singleton_plate_works(i_size, j_size, k_size):
+    def model():
+        with numpyro.plate("i", i_size, dim=-3):
+            with numpyro.plate("j", j_size, dim=-2):
+                with numpyro.plate("k", k_size, dim=-1):
+                    numpyro.sample("a", dist.Normal())
+
+    model = enum(numpyro.handlers.seed(model, rng_seed=0), first_available_dim=-4)
+
+    log_density(model, (), {}, {})