code: def test_zid_shape(gate_shape, base_shape):
      gate = torch.rand(gate_shape)
      base_dist = Normal(
          torch.randn(base_shape), torch.randn(base_shape).exp()
      )
  
      d = ZeroInflatedDistribution(base_dist, gate=gate)
      assert d.batch_shape == broadcast_shape(gate_shape, base_shape)
      assert d.support == base_dist.support
  
      d2 = d.expand([4, 3, 2])
      assert d2.batch_shape == (4, 3, 2)
  
file_path: pyro/tests/distributions/test_zero_inflated.py
imports: import math
  import pytest
  import torch
  from pyro.distributions import Delta, NegativeBinomial, Normal, Poisson, ZeroInflatedDistribution, ZeroInflatedNegativeBinomial, ZeroInflatedPoisson
  from pyro.distributions.util import broadcast_shape
  from tests.common import assert_close
label_desc: Use image node_id 1 for calling a global function with example usage: test_zid_shape(gate_shape, base_shape) without return types