code: def test_branch_merge():
      r"""Test composite function is correctly produced from branching graph.
  
      We would expect the pattern `make_add_sub_mul_pattern` to be merged
      into a single op `add_sub_mul`.
  
         a  b  a  b
          \/    \/
          add  sub                       a  b
           \   /                          \/
            \ /                      add_sub_mul
            mul                     c     |
            /  \                     \    |
         c /  c |       ====>        add_sub_mul
         \/   \/                          |
         add  sub                         |
          \   /                         relu
           \ /
           mul
            |
            |
          relu
      """
  
      pattern_table = [("add_sub_mul", make_add_sub_mul_pattern())]
  
      def before():
          a = relay.var("a", shape=(10, 10))
          b = relay.var("b", shape=(10, 10))
          c = relay.var("c", shape=(10, 10))
          add_node = relay.add(a, b)
          sub_node = relay.subtract(a, b)
          mul_node = relay.multiply(add_node, sub_node)
          add_node_2 = relay.add(c, mul_node)
          sub_node_2 = relay.subtract(c, mul_node)
          mul_node_2 = relay.multiply(add_node_2, sub_node_2)
          r = relay.nn.relu(mul_node_2)
          return relay.Function([a, b, c], r)
  
      def expected():
          a = relay.var("a", shape=(10, 10))
          b = relay.var("b", shape=(10, 10))
          c = relay.var("c", shape=(10, 10))
  
          # add_sub_mul function
          in_1 = relay.var("in_1", shape=(10, 10))
          in_2 = relay.var("in_2", shape=(10, 10))
          add_node = relay.add(in_1, in_2)
          sub_node = relay.subtract(in_1, in_2)
          mul_node = relay.multiply(add_node, sub_node)
          add_sub_mul = relay.Function([in_1, in_2], mul_node)
          add_sub_mul = add_sub_mul.with_attr(
              "Composite", "add_sub_mul"
          )
          add_sub_mul = add_sub_mul.with_attr(
              "PartitionedFromPattern", "add_subtract_multiply_"
          )
  
          # add_sub_mul1 function
          in_3 = relay.var("in_3", shape=(10, 10))
          in_4 = relay.var("in_4", shape=(10, 10))
          add_node_1 = relay.add(in_3, in_4)
          sub_node_1 = relay.subtract(in_3, in_4)
          mul_node_1 = relay.multiply(add_node_1, sub_node_1)
          add_sub_mul_1 = relay.Function([in_3, in_4], mul_node_1)
          add_sub_mul_1 = add_sub_mul_1.with_attr(
              "Composite", "add_sub_mul"
          )
          add_sub_mul_1 = add_sub_mul_1.with_attr(
              "PartitionedFromPattern", "add_subtract_multiply_"
          )
  
          # merged function
          m_add_sub_mul_1 = relay.Call(add_sub_mul, [a, b])
          m_add_sub_mul_2 = relay.Call(
              add_sub_mul_1, [c, m_add_sub_mul_1]
          )
          r = relay.nn.relu(m_add_sub_mul_2)
          return relay.Function([a, b, c], r)
  
      check_result(pattern_table, before(), expected())
  
imports: import pytest
  import tvm
  from tvm import relay, tir
  from tvm.relay.dataflow_pattern import TuplePattern, TupleGetItemPattern, is_op, wildcard
  from tvm.relay.testing import run_opt_pass
label_desc: Use image node_id 9 for calling a global function with example usage: test_branch_merge() and returns: relay, relay