code: def test_combine_parallel_dense():
      """Simple testcase. One dense cannot be combined due to shape mismatch"""
  
      def before(x, w1, w2, w3, w4):
          args = [x, w1, w2, w3, w4]
          y1 = relay.nn.dense(x, w1)
          y2 = relay.nn.dense(x, w2)
  
          # y3 cannot be combined
          y3 = relay.nn.dense(x, w3)
  
          y4 = relay.nn.dense(x, w4)
          y = relay.Tuple((y1, y2, y3, y4))
          return relay.Function(args, y)
  
      def expected(x, w1, w2, w3, w4):
          # use a fixed order of args so alpha equal check can pass
          args = [x, w1, w2, w3, w4]
          x_stacked = relay.stack((x, x, x), axis=0)
          w = relay.stack((w1, w2, w4), axis=0)
          y = relay.nn.batch_matmul(x_stacked, w)
          (y1, y2, y4) = relay.split(y, 3)
          y1 = relay.squeeze(y1, [0])
          y2 = relay.squeeze(y2, [0])
          y4 = relay.squeeze(y4, [0])
  
          # y3 cannot be combined
          y3 = relay.nn.dense(x, w3)
  
          y = relay.Tuple((y1, y2, y3, y4))
          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=(j, k))
          w3 = relay.var("w3", shape=(j + 1, k))
          w4 = relay.var("w4", shape=(j, k))
  
          y_before = before(x, w1, w2, w3, w4)
          y = run_opt_pass(
              y_before,
              transform.CombineParallelDense(min_num_branches=2),
          )
          y_expected = expected(x, w1, w2, w3, w4)
          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 3 for calling a global function with example usage: test_combine_parallel_dense() and returns: relay, relay