code: def test_enum_discrete_plate_shape_broadcasting_ok(
      subsampling, enumerate_
  ):
      @infer.config_enumerate(default=enumerate_)
      def model():
          x_plate = pyro.plate(
              "x_plate",
              5,
              subsample_size=2 if subsampling else None,
              dim=-1,
          )
          y_plate = pyro.plate(
              "y_plate",
              6,
              subsample_size=3 if subsampling else None,
              dim=-2,
          )
          with pyro.plate("num_particles", 50, dim=-3):
              with x_plate:
                  b = pyro.sample(
                      "b",
                      dist.Beta(torch.tensor(1.1), torch.tensor(1.1)),
                  )
              with y_plate:
                  c = pyro.sample("c", dist.Bernoulli(0.5))
              with x_plate, y_plate:
                  d = pyro.sample("d", dist.Bernoulli(b))
  
          # check shapes
          if enumerate_ == "parallel":
              assert b.shape == (50, 1, x_plate.subsample_size)
              assert c.shape == (2, 1, 1, 1)
              assert d.shape == (2, 1, 1, 1, 1)
          elif enumerate_ == "sequential":
              assert b.shape == (50, 1, x_plate.subsample_size)
              assert c.shape in ((), (1, 1, 1))  # both are valid
              assert d.shape in ((), (1, 1, 1))  # both are valid
          else:
              assert b.shape == (50, 1, x_plate.subsample_size)
              assert c.shape == (50, y_plate.subsample_size, 1)
              assert d.shape == (
                  50,
                  y_plate.subsample_size,
                  x_plate.subsample_size,
              )
  
      assert_ok(model, guide=model, max_plate_nesting=3)
  
file_path: pyro/tests/contrib/funsor/test_valid_models_plate.py
imports: import logging
  import pytest
  import torch
  from pyro.ops.indexing import Vindex
  from tests.common import xfail_param
label_desc: Use image node_id 5 for calling a global function with example usage: test_enum_discrete_plate_shape_broadcasting_ok(subsampling, enumerate_) without return types