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
  
file_path: funsor/test/pyro/test_hmm.py
imports: import pyro.distributions
  import pytest
  import torch
  from funsor.pyro.hmm import DiscreteHMM, GaussianHMM, GaussianMRF, SwitchingLinearHMM
  from funsor.testing import assert_close, random_mvn
  from funsor.util import broadcast_shape
label_desc: Use image node_id 1 for calling a global function with example usage: check_expand(old_dist, old_data) without return types