code: def test_script_stacked_rnn(
      seq_len, batch, input_size, hidden_size, num_layers
  ):
      inp = torch.randn(seq_len, batch, input_size)
      states = [
          LSTMState(
              torch.randn(batch, hidden_size),
              torch.randn(batch, hidden_size),
          )
          for _ in range(num_layers)
      ]
      rnn = script_lstm(input_size, hidden_size, num_layers)
      out, out_state = rnn(inp, states)
      custom_state = flatten_states(out_state)
  
      # Control: pytorch native LSTM
      lstm = nn.LSTM(input_size, hidden_size, num_layers)
      lstm_state = flatten_states(states)
      for layer in range(num_layers):
          custom_params = list(rnn.parameters())[
              4 * layer : 4 * (layer + 1)
          ]
          for lstm_param, custom_param in zip(
              lstm.all_weights[layer], custom_params
          ):
              assert lstm_param.shape == custom_param.shape
              with torch.no_grad():
                  lstm_param.copy_(custom_param)
      lstm_out, lstm_out_state = lstm(inp, lstm_state)
  
      assert (out - lstm_out).abs().max() < 1e-5
      assert (custom_state[0] - lstm_out_state[0]).abs().max() < 1e-5
      assert (custom_state[1] - lstm_out_state[1]).abs().max() < 1e-5
  
imports: import numbers
  import warnings
  from collections import namedtuple
  from typing import List, Tuple
  import torch
  import torch.jit
  import torch.nn
  from torch import Tensor
  from torch.nn import Parameter
label_desc: Use image node_id 8 for calling a global function with example usage: test_script_stacked_rnn(seq_len, batch, input_size, hidden_size, num_layers) without return types