code: def test_expand(sample_shape, batch_shape, event_shape):
      ones_shape = torch.Size((1,) * len(batch_shape))
      zero = torch.zeros(ones_shape + event_shape)
      d0 = dist.Uniform(zero - 2, zero + 1).to_event(len(event_shape))
  
      assert d0.sample().shape == ones_shape + event_shape
      assert d0.mean.shape == ones_shape + event_shape
      assert d0.variance.shape == ones_shape + event_shape
      assert (
          d0.sample(sample_shape).shape
          == sample_shape + ones_shape + event_shape
      )
  
      assert (
          d0.expand(sample_shape + batch_shape).batch_shape
          == sample_shape + batch_shape
      )
      assert (
          d0.expand(sample_shape + batch_shape).sample().shape
          == sample_shape + batch_shape + event_shape
      )
      assert (
          d0.expand(sample_shape + batch_shape).mean.shape
          == sample_shape + batch_shape + event_shape
      )
      assert (
          d0.expand(sample_shape + batch_shape).variance.shape
          == sample_shape + batch_shape + event_shape
      )
  
      base_dist = dist.MultivariateNormal(
          torch.zeros(2).expand(*(event_shape + (2,))),
          torch.eye(2).expand(*(event_shape + (2, 2))),
      )
      if len(event_shape) > len(base_dist.batch_shape):
          with pytest.raises(ValueError):
              base_dist.to_event(len(event_shape)).expand(batch_shape)
      else:
          expanded = base_dist.to_event(len(event_shape)).expand(
              batch_shape
          )
          expanded_batch_ndims = getattr(
              expanded, "reinterpreted_batch_ndims", 0
          )
          assert expanded.batch_shape == batch_shape
          assert expanded.event_shape == (
              base_dist.batch_shape[
                  len(base_dist.batch_shape) - expanded_batch_ndims :
              ]
              + base_dist.event_shape
          )
  
file_path: pyro/tests/distributions/test_independent.py
imports: import pytest
  import torch
  from torch.distributions.utils import _sum_rightmost
  import pyro.distributions
  from pyro.util import torch_isnan
  from tests.common import assert_equal
label_desc: Use image node_id 3 for calling a global function with example usage: test_expand(sample_shape, batch_shape, event_shape) without return types