Repository: pyro-ppl/numpyro

Issue:
Declare dependent dimensions
In Pyro, we can declare dependent dimensions through `.to_event()`. Currently, this is useful for correctly computing log_prob of a transformed distribution where `transform.event_dim>0` but `base_dist.event_shape=()`. In particular, we need this to declare `base_dist` in autoguide.

In Pyro, this is achieved by using Independent class. It seems unnecessary to me. Can we just manipulate `batch_shape`, `event_shape` and correctly compute the returned value for `log_prob`? I think we just need 1 line of code for each distribution
```
def log_prob(...):
    ...
    # event_ndim minus 1 for multivariate distributions
    return sum_rightmost(log_prob, len(self.event_shape))
```

If Independent class turns out to be useful, then we will implement it instead.


Discussion:
I think following the Pyro pattern here will be more modular. Adjusting the `batch_shape`, `event_shape` inside each distribution will lead to issues like `batch_shape` not matching the shape of distribution parameters which will necessitate all kinds of checks that are avoidable. Using `.to_event` will confine the change to a single class / method. I am happy to implement a version of this.
Sure, I didn't follow up with Distribution development from the start so I didn't appreciate much the current design decision. I imagine the method `to_event` can be implemented as
```
def to_event(self, n):
    self.event_shape = self.batch_shape[-n:] + self.event_shape
    self.batch_shape = self.batch_shape[:-n]
```
For some situations such as `loc.shape=(10, 1)`, `scale.shape=(1, 10)`, this still works.

While writing this reply, I do agree that using Independent will be beneficial for functions which utilizes this dependent information. It will take more effort to check how many dependent dims a distribution has (because some are univariate, some are multivariate) if we don't use it.  **Edit**: hmm, we can just add `event_dim = 0/1/2` for each distribution (depends on it is univariate or multivariate) and compute `dependent_ndim = len(event_shape) - event_dim`...
Np, it is always good to question existing design patterns! Regarding the solution above, it is not that it cannot be done, but this will leak into individual distribution classes (methods like `log_prob`; additionally, if any methods use batch shape in the constructor to do something like set some internal attribute, and we later change batch shape, we will need to also change what was done internally in the constructor). 

This also makes it harder for other user-defined distributions that implement the `Distribution` interface to adhere to because all distribution methods and attributes need to assume that `batch_shape` and `event_shape` are mutable attributes. Therefore any method that makes use of these attributes to set some state may need to be rerun when these are changed. On the other hand, there are no specific design advantages to doing this, as far as I can see. Those are some of the reasons why this was implemented as a composition over an existing distribution instance in both pytorch and tensorflow.
Thanks @neerajprad ! Based on your comment, it looks like an effect handler to me. :D

Resolving patch:
diff --git a/numpyro/distributions/__init__.py b/numpyro/distributions/__init__.py
--- a/numpyro/distributions/__init__.py
+++ b/numpyro/distributions/__init__.py
@@ -36,7 +36,7 @@
     Poisson,
     PRNGIdentity
 )
-from numpyro.distributions.distribution import Distribution, TransformedDistribution
+from numpyro.distributions.distribution import Distribution, Independent, TransformedDistribution
 
 __all__ = [
     'Bernoulli',
@@ -59,6 +59,7 @@
     'GaussianRandomWalk',
     'HalfCauchy',
     'HalfNormal',
+    'Independent',
     'InverseGamma',
     'LKJCholesky',
     'LogNormal',
diff --git a/numpyro/distributions/distribution.py b/numpyro/distributions/distribution.py
--- a/numpyro/distributions/distribution.py
+++ b/numpyro/distributions/distribution.py
@@ -166,6 +166,75 @@ def __call__(self, *args, **kwargs):
             return self.sample_with_intermediates(key, *args, **kwargs)
         return self.sample(key, *args, **kwargs)
 
+    def to_event(self, reinterpreted_batch_ndims=None):
+        if reinterpreted_batch_ndims is None:
+            reinterpreted_batch_ndims = len(self.batch_shape)
+        return Independent(self, reinterpreted_batch_ndims)
+
+
+class Independent(Distribution):
+    """
+    Reinterprets batch dimensions of a distribution as event dims by shifting
+    the batch-event dim boundary further to the left.
+
+    From a practical standpoint, this is useful when changing the result of
+    :meth:`log_prob`. For example, a univariate Normal distribution can be
+    interpreted as a multivariate Normal with diagonal covariance::
+
+    .. testsetup::
+
+       import numpyro.distributions as dist
+
+    .. doctest::
+
+        >>> normal = dist.Normal(np.zeros(3), np.ones(3))
+        >>> [normal.batch_shape, normal.event_shape]
+        [torch.Size((3,)), torch.Size(())]
+        >>> diag_normal = Independent(normal, 1)
+        >>> [diag_normal.batch_shape, diag_normal.event_shape]
+        [torch.Size(()), torch.Size((3,))]
+
+    :param numpyro.distribution.Distribution base_distribution: a distribution instance.
+    :param int reinterpreted_batch_ndims: the number of batch dims to reinterpret as event dims.
+    """
+    arg_constraints = {}
+
+    def __init__(self, base_dist, reinterpreted_batch_ndims, validate_args=None):
+        if reinterpreted_batch_ndims > len(base_dist.batch_shape):
+            raise ValueError("Expected reinterpreted_batch_ndims <= len(base_distribution.batch_shape), "
+                             "actual {} vs {}".format(reinterpreted_batch_ndims,
+                                                      len(base_dist.batch_shape)))
+        shape = base_dist.batch_shape + base_dist.event_shape
+        event_dim = reinterpreted_batch_ndims + len(base_dist.event_shape)
+        batch_shape = shape[:len(shape) - event_dim]
+        event_shape = shape[len(shape) - event_dim:]
+        self.base_dist = base_dist
+        self.reinterpreted_batch_ndims = reinterpreted_batch_ndims
+        super(Independent, self).__init__(batch_shape, event_shape, validate_args=validate_args)
+
+    @property
+    def support(self):
+        return self.base_dist.support
+
+    @property
+    def reparameterized_params(self):
+        return self.base_dist.reparameterized_params
+
+    @property
+    def mean(self):
+        return self.base_dist.mean
+
+    @property
+    def variance(self):
+        return self.base_dist.variance
+
+    def sample(self, key, sample_shape=()):
+        return self.base_dist.sample(key, sample_shape=sample_shape)
+
+    def log_prob(self, value):
+        log_prob = self.base_dist.log_prob(value)
+        return sum_rightmost(log_prob, self.reinterpreted_batch_ndims)
+
 
 class TransformedDistribution(Distribution):
     """


Test patch:
diff --git a/test/test_distributions.py b/test/test_distributions.py
--- a/test/test_distributions.py
+++ b/test/test_distributions.py
@@ -355,6 +355,20 @@ def test_log_prob(jax_dist, sp_dist, params, prepend_shape, jit):
     assert_allclose(jit_fn(jax_dist.log_prob)(samples), expected, atol=1e-5)
 
 
+@pytest.mark.parametrize('jax_dist, sp_dist, params', CONTINUOUS + DISCRETE)
+def test_independent_shape(jax_dist, sp_dist, params):
+    d = jax_dist(*params)
+    batch_shape, event_shape = d.batch_shape, d.event_shape
+    shape = batch_shape + event_shape
+    for i in range(len(batch_shape)):
+        indep = dist.Independent(d, reinterpreted_batch_ndims=i)
+        sample = indep.sample(random.PRNGKey(0))
+        event_boundary = len(shape) - len(event_shape) - i
+        assert indep.batch_shape == shape[:event_boundary]
+        assert indep.event_shape == shape[event_boundary:]
+        assert np.shape(indep.log_prob(sample)) == shape[:event_boundary]
+
+
 def _tril_cholesky_to_tril_corr(x):
     w = vec_to_tril_matrix(x, diagonal=-1)
     diag = np.sqrt(1 - np.sum(w ** 2, axis=-1))