code: def test_expand(sample_shape, batch_shape, event_shape):
      ones_shape = torch.Size((1,) * len(batch_shape))
      mask = torch.empty(ones_shape).bernoulli_(0.5).bool()
      zero = torch.zeros(ones_shape + event_shape)
      d0 = dist.Uniform(zero - 2, zero + 1).to_event(len(event_shape))
      d1 = dist.Uniform(zero - 1, zero + 2).to_event(len(event_shape))
      d = dist.MaskedMixture(mask, d0, d1)
  
      assert d.sample().shape == ones_shape + event_shape
      assert d.mean.shape == ones_shape + event_shape
      assert d.variance.shape == ones_shape + event_shape
      assert (
          d.sample(sample_shape).shape
          == sample_shape + ones_shape + event_shape
      )
  
      assert (
          d.expand(sample_shape + batch_shape).batch_shape
          == sample_shape + batch_shape
      )
      assert (
          d.expand(sample_shape + batch_shape).sample().shape
          == sample_shape + batch_shape + event_shape
      )
      assert (
          d.expand(sample_shape + batch_shape).mean.shape
          == sample_shape + batch_shape + event_shape
      )
      assert (
          d.expand(sample_shape + batch_shape).variance.shape
          == sample_shape + batch_shape + event_shape
      )
  
file_path: pyro/tests/distributions/test_mixture.py
imports: import pytest
  import torch
  import pyro.distributions
  from pyro.distributions.util import broadcast_shape
  from pyro.util import torch_isnan
  from tests.common import assert_equal
label_desc: Use image node_id 4 for calling a global function with example usage: test_expand(sample_shape, batch_shape, event_shape) without return types