code: def test_expand(sample_shape, batch_shape):
      log_factor = torch.randn(batch_shape)
      d1 = dist.Unit(log_factor)
      v1 = d1.sample()
  
      d2 = d1.expand(sample_shape + batch_shape)
      assert d2.batch_shape == sample_shape + batch_shape
      v2 = d2.sample()
      assert v2.shape == sample_shape + batch_shape + (0,)
      assert_equal(d1.log_prob(v2), d2.log_prob(v1))
  
imports: import pytest
  import torch
  import pyro.distributions
  from tests.common import assert_equal
label_desc: Use image node_id 2 for calling a global function with example usage: test_expand(sample_shape, batch_shape) without return types