code: def einsum(equation, *operands):
      """
      Forward-max-sum backward-argmax implementation of einsum.
      This assumes all operands have a ``._pyro_dims`` attribute set.
      """
      inputs, output = equation.split("->")
      inputs = inputs.split(",")
  
      contract_dims = "".join(
          sorted(set().union(*inputs) - set(output))
      )
      dims = output + contract_dims
      result = reduce(
          operator.add,
          broadcast_all(*operands, inputs=inputs, dims=dims),
      )
      if contract_dims:
          output_shape = result.shape[: len(output)]
          result = ops.amax(result.reshape(output_shape + (-1,)), -1)
      elif result is operands[0]:
          result = result[...]  # create a new object
      assert len(result.shape) == len(output)
  
      return result
  
imports: import operator
  from functools import reduce
  import funsor.ops
  from funsor.einsum.util import Tensordot, broadcast_all
label_desc: Use image node_id 1 for calling a global function with example usage: einsum(equation) and returns: result