code: def test_shape(backend, equation):
      backend = "pyro.ops.einsum.torch_{}".format(backend)
      inputs, output = equation.split("->")
      inputs = inputs.split(",")
      symbols = sorted(set(equation) - set(",->"))
      sizes = dict(zip(symbols, itertools.count(2)))
      input_shapes = [
          torch.Size(sizes[dim] for dim in dims) for dims in inputs
      ]
      operands = [torch.randn(shape) for shape in input_shapes]
      for input_, x in zip(inputs, operands):
          x._pyro_dims = input_
  
      # check forward pass
      for x in operands:
          require_backward(x)
      expected = contract(
          equation, *operands, backend="pyro.ops.einsum.torch_log"
      )
      actual = contract(equation, *operands, backend=backend)
      if backend.endswith("map"):
          assert actual.dtype == expected.dtype
          assert actual.shape == expected.shape
      else:
          assert_equal(actual, expected)
  
      # check backward pass
      actual._pyro_backward()
      for input_, x in zip(inputs, operands):
          backward_result = x._pyro_backward_result
          if backend.endswith("marginal"):
              assert backward_result.shape == x.shape
          else:
              contract_dims = set(input_) - set(output)
              if contract_dims:
                  assert backward_result.size(0) == len(contract_dims)
                  assert set(backward_result._pyro_dims[1:]) == set(
                      output
                  )
                  for sample, dim in zip(
                      backward_result, backward_result._pyro_sample_dims
                  ):
                      assert sample.min() >= 0
                      assert sample.max() < sizes[dim]
              else:
                  assert backward_result is None
  
file_path: pyro/tests/ops/einsum/test_adjoint.py
imports: import gc
  import itertools
  import pytest
  import torch
  from pyro.ops.einsum import contract
  from pyro.ops.einsum.adjoint import require_backward
  from tests.common import assert_equal
label_desc: Use image node_id 1 for calling a global function with example usage: test_shape(backend, equation) without return types