code: def test_custom_lstm():
      input_name = "input"
      states_name = "states"
      seq_len = 5
      batch = 2
      input_size = 3
      hidden_size = 4
      num_layers = 3
      state_tensor_shape = (batch, hidden_size)
  
      torch.manual_seed(1)
  
      inp = torch.randn(seq_len, batch, input_size)
  
      input_shapes = [
          (input_name, (seq_len, batch, input_size)),
          (states_name, (state_tensor_shape, state_tensor_shape)),
      ]
  
      input_shapes_stacked = [
          (input_name, (seq_len, batch, input_size)),
          (
              states_name,
              [
                  (state_tensor_shape, state_tensor_shape),
                  (state_tensor_shape, state_tensor_shape),
              ],
          ),
      ]
  
      input_shapes_stacked_bidir = [
          (input_name, (seq_len, batch, input_size)),
          (
              states_name,
              [
                  [
                      (state_tensor_shape, state_tensor_shape)
                      for _ in range(2)
                  ]
                  for _ in range(num_layers)
              ],
          ),
      ]
  
      states = [
          (
              torch.randn(state_tensor_shape),
              torch.randn(state_tensor_shape),
          )
          for _ in range(num_layers)
      ]
  
      bidir_states = [
          (
              torch.randn(state_tensor_shape),
              torch.randn(state_tensor_shape),
          )
          for _ in range(2)
      ]
  
      stacked_bidir_states = [
          [
              (
                  torch.randn(state_tensor_shape),
                  torch.randn(state_tensor_shape),
              )
              for _ in range(2)
          ]
          for _ in range(num_layers)
      ]
  
      models = [
          (
              "lstm",
              lstm(input_size, hidden_size).eval(),
              states[0],
              input_shapes,
          ),
          (
              "stacked",
              stacked_lstm(input_size, hidden_size, num_layers).eval(),
              states,
              input_shapes_stacked,
          ),
          (
              "bidir",
              bidir_lstm(input_size, hidden_size).eval(),
              bidir_states,
              input_shapes_stacked,
          ),
          # TODO(masahi): stacked bidir seems to have a rare accuracy issue
          # (
          #     "stacked_bidir",
          #     stacked_bidir_lstm(input_size, hidden_size, num_layers).eval(),
          #     stacked_bidir_states,
          #     input_shapes_stacked_bidir,
          # ),
      ]
  
      for name, raw_model, states, input_shapes in models:
          script_module = torch.jit.script(raw_model)
          with tvm.testing.disable_span_filling():
              mod, params = from_pytorch(script_module, input_shapes)
          with tvm.testing.enable_span_filling():
              mod_with_span, _ = from_pytorch(
                  script_module, input_shapes
              )
          assert tvm.ir.structural_equal(
              mod, mod_with_span, map_free_vars=True
          )
  
          with torch.no_grad():
              pt_result = raw_model(inp.clone(), states)
  
          params[input_name] = inp.numpy()
  
          if isinstance(states, tuple):
              states_np = tuple(st.numpy() for st in states)
          elif isinstance(states, list) and isinstance(
              states[0], torch.Tensor
          ):
              states_np = [st.numpy() for st in states]
          elif isinstance(states, list) and isinstance(
              states[0], tuple
          ):
              states_np = [
                  tuple(st.numpy() for st in states[i])
                  for i in range(len(states))
              ]
          elif isinstance(states, list) and isinstance(states[0], list):
              states_np = [
                  [
                      tuple(st.numpy() for st in states)
                      for states in states[layer]
                  ]
                  for layer in range(num_layers)
              ]
          else:
              assert False
  
          if isinstance(states_np, list):
              params[states_name] = convert_list_to_vmobj(states_np)
          else:
              params[states_name] = states_np
  
          for tgt, dev in tvm.testing.enabled_targets():
              print("Running %s on target %s" % (name, tgt))
              run_and_compare(
                  mod, params, pt_result, target=tgt, device=dev
              )
  
imports: import numpy
  import torch
  import torch.nn
  from torch.nn import Parameter
  import torch.jit
  from typing import List, Tuple
  from torch import Tensor
  import tvm
  import tvm.testing
  from tvm import relay
  from tvm.relay.frontend.pytorch import from_pytorch
  from tvm.relay.prelude import Prelude
  from tvm.runtime.container import ADT, tuple_object
label_desc: Use image node_id 10 for calling a global function with example usage: test_custom_lstm() without return types