code: def test_legalize():
      """Test directly replacing an operator with a new one"""
  
      def before():
          x = relay.var("x", shape=(1, 64, 56, 56))
          weight = relay.var("weight", shape=(64, 64, 3, 3))
          y = relay.nn.conv2d(
              x, weight, channels=64, kernel_size=(3, 3), padding=(1, 1)
          )
          y = relay.nn.relu(y)
          y = relay.Function([x, weight], y)
          return y
  
      def legalize_conv2d(attrs, inputs, types):
          data, weight = inputs
          weight = relay.multiply(weight, relay.const(2.0, "float32"))
          return relay.nn.conv2d(data, weight, **attrs)
  
      def expected():
          x = relay.var("x", shape=(1, 64, 56, 56))
          weight = relay.var("weight", shape=(64, 64, 3, 3))
          y = relay.nn.conv2d(
              x,
              relay.multiply(weight, relay.const(2.0, "float32")),
              channels=64,
              kernel_size=(3, 3),
              padding=(1, 1),
          )
          y = relay.nn.relu(y)
          y = relay.Function([x, weight], y)
          return y
  
      with TempOpAttr("nn.conv2d", "FTVMLegalize", legalize_conv2d):
          a = before()
          a = run_opt_pass(a, transform.Legalize())
          b = run_opt_pass(expected(), transform.InferType())
  
      assert tvm.ir.structural_equal(a, b), "Actual = \n" + str(a)
  
file_path: tvm/tests/python/relay/test_pass_legalize.py
imports: import numpy
  import pytest
  import tvm
  from tvm import te
  from tvm import relay
  from tvm.contrib import graph_executor
  from tvm.relay import transform, analysis
  from tvm.relay.testing.temp_op_attr import TempOpAttr
label_desc: Use image node_id 2 for calling a global function with example usage: test_legalize() and returns: y, relay, y