code: def check_rnn(
      rnn_type,
      rnn_mod,
      target=tvm.target.Target("llvm -mcpu=core-avx2"),
      dev=tvm.cpu(0),
  ):
      def get_model(
          rnn_type,
          rnn_mod,
          args,
      ):
          # Fill args
          if "b" in rnn_mod:
              args["bidirectional"] = True
          if "s" in rnn_mod:
              args["layer_num"] = num_layers
          if "tanh" in rnn_mod:
              args["nonlinearity"] = "tanh"
          if "relu" in rnn_mod:
              args["nonlinearity"] = "relu"
  
          if rnn_type == "GRU":
              RNN_Model_selector = GRU_Model
          elif rnn_type == "LSTM":
              RNN_Model_selector = LSTM_Model
              if "p" in rnn_mod:
                  args["proj_size"] = lstm_projection_size
          elif rnn_type == "RNN":
              RNN_Model_selector = RNN_Model_Impl
  
          return RNN_Model_selector(**args)
  
      def get_onnx_model(model):
          onnx_io = io.BytesIO()
          with torch.no_grad():
              input_names = model.get_input_names()
              inputs = model.get_dummy_inputs()
  
              # default export (without dynamic input)
              torch.onnx.export(
                  model, inputs, onnx_io, input_names=input_names
              )
  
          onnx_io.seek(0, 0)
          return onnx.load_model(onnx_io)
  
      model = None
      dtype = "float32"
      device = torch.device("cpu")
      for batch_first in (True, False):
          for use_bias in (True, False):
              for rnd_weights in [True]:  # (True, False):
                  model_inputs = {
                      "batch_first": batch_first,
                      "use_bias": use_bias,
                      "rnd_weights_init": rnd_weights,
                  }
                  model = get_model(rnn_type, rnn_mod, model_inputs)
                  model.to(device)
                  model.eval()
  
                  # Get golden output from original model
                  dummy_inputs = model.get_dummy_inputs()
                  golden_output = (
                      model.forward(dummy_inputs[0].to(device))
                      .detach()
                      .cpu()
                      .numpy()
                  )
  
                  tvm_output = None
                  for format in ["pt"]:  # ["pt", "onnx"]:
                      shape_desc = model.get_shape_desc(format)
                      if format == "pt":
                          # Use torch.jit.trace to generate a torch.jit.ScriptModule via tracing.
                          traced_script_module = torch.jit.trace(
                              model, dummy_inputs[0]
                          ).eval()
  
                          # Import model to Relay
                          with tvm.testing.disable_span_filling():
                              mod, params = relay.frontend.from_pytorch(
                                  traced_script_module, shape_desc
                              )
                          with tvm.testing.enable_span_filling():
                              (
                                  mod_with_span,
                                  _,
                              ) = relay.frontend.from_pytorch(
                                  traced_script_module, shape_desc
                              )
                          assert tvm.ir.structural_equal(
                              mod, mod_with_span, map_free_vars=True
                          )
                      elif format == "onnx":
                          try:
                              onnx_model = get_onnx_model(model)
                          except:
                              print(
                                  "WARNING: torch.onnx.export does not support conversion LSTM with projection "
                                  "from pytorch! TODO: waiting for the support and correct test after that."
                              )
                              continue
  
                          # Import model to Relay
                          with tvm.testing.disable_span_filling():
                              mod, params = relay.frontend.from_onnx(
                                  onnx_model, shape_desc
                              )
                          with tvm.testing.enable_span_filling():
                              (
                                  mod_with_span,
                                  _,
                              ) = relay.frontend.from_onnx(
                                  onnx_model, shape_desc
                              )
                          assert tvm.ir.structural_equal(
                              mod, mod_with_span, map_free_vars=True
                          )
  
                      # Model compilation by tvm
                      with tvm.transform.PassContext(opt_level=3):
                          lib = relay.build(
                              mod, target=target, params=params
                          )
  
                      # Inference of the model with given input data
                      m = graph_executor.GraphModule(
                          lib["default"](dev)
                      )
  
                      # Set inputs
                      tvm_inputs = model.get_tvm_inputs(dtype)
                      m.set_input(**tvm_inputs)
                      # Execute
                      m.run()
                      # Get outputs (converted to numpy array)
                      tvm_output = m.get_output(0).numpy()
  
                      compare(tvm_output, golden_output)
  
imports: import torch
  import tvm
  import tvm.testing
  import onnx
  import io
  import sys
  from tvm import relay
  from tvm.contrib import graph_executor
  from torch import nn
label_desc: Use image node_id 3 for calling a global function with example usage: check_rnn(rnn_type, rnn_mod, target, dev) and returns: RNN_Model_selector, onnx