code: def test_einsum_categorical(equation):
      if get_backend() == "jax":
          from funsor.jax.distributions import Categorical
      else:
          from funsor.torch.distributions import Categorical
  
      inputs, outputs, sizes, operands, _ = make_einsum_example(
          equation
      )
      operands = [
          ops.abs(operand) / ops.abs(operand).sum(-1)[..., None]
          for operand in operands
      ]
  
      expected = opt_einsum.contract(
          equation,
          *operands,
          backend=BACKEND_TO_EINSUM_BACKEND[get_backend()],
      )
  
      with reflect:
          funsor_operands = [
              Categorical(
                  probs=Tensor(
                      operand,
                      inputs=OrderedDict(
                          [(d, Bint[sizes[d]]) for d in inp[:-1]]
                      ),
                  )
              )(value=Variable(inp[-1], Bint[sizes[inp[-1]]])).exp()
              for inp, operand in zip(inputs, operands)
          ]
  
          naive_ast = naive_einsum(equation, *funsor_operands)
          optimized_ast = apply_optimizer(naive_ast)
  
      print("Naive expression: {}".format(naive_ast))
      print("Optimized expression: {}".format(optimized_ast))
      actual_optimized = reinterpret(optimized_ast)  # eager by default
      actual = naive_einsum(
          equation, *map(reinterpret, funsor_operands)
      )
  
      if len(outputs[0]) > 0:
          actual = actual.align(tuple(outputs[0]))
          actual_optimized = actual_optimized.align(tuple(outputs[0]))
  
      assert_close(actual, actual_optimized, atol=1e-4)
  
      assert expected.shape == actual.data.shape
      assert_close(expected, actual.data)
      for output in outputs:
          for i, output_dim in enumerate(output):
              assert output_dim in actual.inputs
              assert (
                  actual.inputs[output_dim].dtype == sizes[output_dim]
              )
  
file_path: funsor/test/test_einsum.py
imports: from collections import OrderedDict
  import opt_einsum
  import pytest
  import funsor
  import funsor.ops
  from funsor.cnf import BACKEND_TO_EINSUM_BACKEND, BACKEND_TO_LOGSUMEXP_BACKEND, BACKEND_TO_MAP_BACKEND
  from funsor.domains import Bint
  from funsor.einsum import naive_einsum, naive_plated_einsum
  from funsor.interpretations import reflect
  from funsor.interpreter import reinterpret
  from funsor.optimizer import apply_optimizer
  from funsor.tensor import Tensor
  from funsor.terms import Variable
  from funsor.testing import assert_close, excludes_backend, make_einsum_example, requires_backend
  from funsor.util import get_backend
label_desc: Use image node_id 3 for calling a global function with example usage: test_einsum_categorical(equation) without return types