code: def test_combine_parallel_dense_flat():
      """Simple testcase. All matmul of different output dim can be combined"""
  
      def before(x, w1, w2, w3):
          args = [x, w1, w2, w3]
          y1 = relay.nn.dense(x, w1)
          y2 = relay.nn.dense(x, w2)
          y3 = relay.nn.dense(x, w3)
          y = relay.Tuple((y1, y2, y3))
          return relay.Function(args, y)
  
      def expected(x, w1, w2, w3, j):
          args = [x, w1, w2, w3]
          w_stacked = relay.concatenate((w1, w2, w3), axis=0)
          y = relay.nn.dense(x, w_stacked, units=6 * j)
          strides = [1, 1]
          y1 = relay.strided_slice(
              y,
              begin=[0, 0],
              end=[-1, j],
              strides=strides,
              slice_mode="size",
          )
          y2 = relay.strided_slice(
              y,
              begin=[0, j],
              end=[-1, 2 * j],
              strides=strides,
              slice_mode="size",
          )
          y3 = relay.strided_slice(
              y,
              begin=[0, 3 * j],
              end=[-1, 3 * j],
              strides=strides,
              slice_mode="size",
          )
          y = relay.Tuple((y1, y2, y3))
          return relay.Function(args, y)
  
      def check(i, j, k):
          x = relay.var("x", shape=(i, k))
          w1 = relay.var("w1", shape=(j, k))
          w2 = relay.var("w2", shape=(2 * j, k))
          w3 = relay.var("w3", shape=(3 * j, k))
  
          y_before = before(x, w1, w2, w3)
          combine_pass = transform.CombineParallelDense(
              min_num_branches=3, to_batch=False
          )
          y = run_opt_pass(y_before, combine_pass)
          y_expected = expected(x, w1, w2, w3, j)
          y_expected = run_opt_pass(y_expected, transform.InferType())
          tvm.ir.assert_structural_equal(
              y, y_expected, map_free_vars=True
          )
  
      check(3, 5, 4)
      check(100, 200, 300)
  
file_path: tvm/tests/python/relay/test_pass_combine_parallel_dense.py
imports: import tvm
  import tvm.testing
  from tvm import relay
  from tvm.relay import transform
label_desc: Use image node_id 6 for calling a global function with example usage: test_combine_parallel_dense_flat() and returns: relay, relay