code: def check_expand(old_dist, old_data):
      new_batch_shape = (2,) + old_dist.batch_shape
      new_dist = old_dist.expand(new_batch_shape)
      assert new_dist.batch_shape == new_batch_shape
  
      old_log_prob = new_dist.log_prob(old_data)
      assert old_log_prob.shape == new_batch_shape
  
      new_data = old_data.expand(new_batch_shape + new_dist.event_shape)
      new_log_prob = new_dist.log_prob(new_data)
      assert_close(old_log_prob, new_log_prob)
      assert new_dist.log_prob(new_data).shape == new_batch_shape
  
imports: import logging
  import operator
  from functools import reduce
  import opt_einsum
  import pytest
  import torch
  import pyro
  import pyro.distributions
  from pyro.distributions.hmm import _sequential_gamma_gaussian_tensordot, _sequential_logmatmulexp
  from pyro.distributions.util import broadcast_shape
  from pyro.infer import TraceEnum_ELBO, config_enumerate
  from pyro.ops.gamma_gaussian import gamma_and_mvn_to_gamma_gaussian, gamma_gaussian_tensordot, matrix_and_mvn_to_gamma_gaussian
  from pyro.ops.gaussian import gaussian_tensordot, matrix_and_mvn_to_gaussian, mvn_to_gaussian
  from pyro.ops.indexing import Vindex
  from tests.common import assert_close
  from tests.ops.gamma_gaussian import assert_close_gamma_gaussian, random_gamma, random_gamma_gaussian
  from tests.ops.gaussian import random_mvn
label_desc: Use image node_id 1 for calling a global function with example usage: check_expand(old_dist, old_data) without return types