code: def test_lstmp():
      hidden_size, projection_size = 512, 256
      rtol, atol = 1e-4, 1e-4
      batch_size, seq_len = 5, 3
      input_size = 128
      lstm_input = mx.np.random.uniform(
          size=(seq_len, batch_size, input_size)
      )
  
      # ==== Unidirectional Layer ====
      for num_layers in [1, 3]:
          fused_layer = gluon.rnn.LSTM(
              hidden_size,
              projection_size=projection_size,
              num_layers=num_layers,
              layout="TNC",
              bidirectional=False,
          )
  
          stack_layer = mx.gluon.rnn.HybridSequentialRNNCell()
          for _ in range(num_layers):
              stack_layer.add(
                  gluon.rnn.LSTMPCell(
                      hidden_size, projection_size=projection_size
                  )
              )
          fused_layer.initialize()
  
          fused_begin_state = fused_layer.begin_state(batch_size)
          stack_begin_state = stack_layer.begin_state(
              batch_size=batch_size
          )
          fused_layer.infer_shape(lstm_input, fused_begin_state)
          stack_layer.infer_shape(0, lstm_input, False)
          stack_layer.initialize()
          fused_layer_params = fused_layer.collect_params()
          stack_layer_params = stack_layer.collect_params()
  
          fused_weight_shape = fused_layer_params["rnn_param"].shape
          print(fused_weight_shape)
          w = mx.np.zeros(shape=fused_weight_shape)
          fused_layer_params["rnn_param"].set_data(w)
          fused_layer_params_split = split_rnn_params(
              w,
              "lstm",
              num_layers,
              input_size,
              hidden_size,
              False,
              projection_size=projection_size,
          )
  
          for name, value in fused_layer_params_split.items():
              stack_layer_params[
                  name[1:].replace("_", ".", 1)
              ].set_data(value)
  
          fused_output, fused_states = fused_layer(
              lstm_input.copy(), fused_begin_state
          )
          stack_output, stack_states = stack_layer.unroll(
              seq_len,
              lstm_input.copy(),
              begin_state=stack_begin_state,
              layout="TNC",
              merge_outputs=True,
          )
  
          assert_almost_equal(
              fused_output.asnumpy(),
              stack_output.asnumpy(),
              rtol=rtol,
              atol=atol,
          )
          check_rnn_states(
              fused_states, stack_states, num_layers, False
          )
  
      # ==== Bidirectional Layer ====
      for num_layers in [1, 3]:
          fused_layer = gluon.rnn.LSTM(
              hidden_size,
              projection_size=projection_size,
              num_layers=num_layers,
              layout="TNC",
              bidirectional=True,
          )
  
          stack_layer = mx.gluon.rnn.HybridSequentialRNNCell()
          for _ in range(num_layers):
              stack_layer.add(
                  gluon.rnn.BidirectionalCell(
                      gluon.rnn.LSTMPCell(
                          hidden_size, projection_size=projection_size
                      ),
                      gluon.rnn.LSTMPCell(
                          hidden_size, projection_size=projection_size
                      ),
                  )
              )
          fused_layer.initialize()
          stack_layer.initialize()
  
          fused_begin_state = fused_layer.begin_state(batch_size)
          stack_begin_state = stack_layer.begin_state(
              batch_size=batch_size
          )
          fused_layer.infer_shape(lstm_input, fused_begin_state)
          stack_layer.infer_shape(0, lstm_input, False)
          stack_layer.initialize()
          fused_layer_params = fused_layer.collect_params()
          stack_layer_params = stack_layer.collect_params()
  
          fused_weight_shape = fused_layer_params["rnn_param"].shape
          w = mx.np.zeros(shape=fused_weight_shape)
          fused_layer_params["rnn_param"].set_data(w)
          fused_layer_params_split = split_rnn_params(
              w,
              "lstm",
              num_layers,
              input_size,
              hidden_size,
              True,
              projection_size=projection_size,
          )
  
          for name, value in fused_layer_params_split.items():
              cur = name.split("_")[0]
              stack_layer_params[
                  "{}.{}_cell.{}".format(
                      cur[1:], name[0], name[len(cur) + 1 :]
                  )
              ].set_data(value)
  
          fused_output, fused_states = fused_layer(
              lstm_input.copy(), fused_begin_state
          )
          stack_output, stack_states = stack_layer.unroll(
              seq_len,
              lstm_input.copy(),
              begin_state=stack_begin_state,
              layout="TNC",
              merge_outputs=True,
          )
  
          assert_almost_equal(
              fused_output.asnumpy(),
              stack_output.asnumpy(),
              rtol=rtol,
              atol=atol,
          )
          check_rnn_states(fused_states, stack_states, num_layers, True)
  
imports: import mxnet
  from mxnet import gluon, np
  import numpy
  import copy
  from functools import partial
  from numpy.testing import assert_allclose
  from mxnet.gluon.utils import split_rnn_params
  import pytest
  from mxnet.test_utils import almost_equal, assert_almost_equal, default_device
  from common import assert_raises_cudnn_not_satisfied, retry
label_desc: Use image node_id 4 for calling a global function with example usage: test_lstmp() without return types