code: def test_break_plates():
      shape = torch.Size([5, 4, 3, 2])
      x = torch.arange(shape.numel()).reshape(shape)
  
      MockPlate = namedtuple("MockPlate", "dim, size")
      h = MockPlate(-4, 6)
      i = MockPlate(-3, 5)
      j = MockPlate(-2, 4)
      k = MockPlate(-1, 3)
  
      actual = _break_plates(x, {i, j, k}, set())
      expected = x.reshape(-1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {i})
      expected = x.reshape(5, 1, 1, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {j})
      expected = x.permute((1, 0, 2, 3)).reshape(4, 1, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {k})
      expected = x.permute((2, 0, 1, 3)).reshape(3, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {i, j})
      expected = x.reshape(5, 4, 1, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {i, k})
      expected = x.permute((0, 2, 1, 3)).reshape(5, 1, 3, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {j, k})
      expected = x.permute((1, 2, 0, 3)).reshape(4, 3, -1)
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {i, j, k})
      expected = x
      assert_equal(actual, expected)
  
      actual = _break_plates(x, {i, j, k}, {h, i, j, k})
      expected = x
      assert_equal(actual, expected)
  
imports: from collections import OrderedDict, namedtuple
  import pytest
  import torch
  import pyro
  import pyro.distributions
  import pyro.poutine
  from pyro.infer import SVI, JitTrace_ELBO, Predictive, Trace_ELBO
  from pyro.infer.autoguide import AutoGaussian, AutoGuideList
  from pyro.infer.autoguide.gaussian import AutoGaussianDense, AutoGaussianFunsor, _break_plates
  from pyro.infer.reparam import LocScaleReparam
  from pyro.optim import ClippedAdam
  from tests.common import assert_close, assert_equal, xfail_if_not_implemented
label_desc: Use image node_id 1 for calling a global function with example usage: test_break_plates() without return types