code: def check_rnn_states(
      fused_states,
      stack_states,
      num_layers,
      bidirectional=False,
      is_lstm=True,
  ):
      directions = 2 if bidirectional else 1
      assert (
          len(stack_states) / len(fused_states)
          == num_layers * directions
      )
  
      fused_states = [state.asnumpy() for state in fused_states]
      stack_states = [
          _np.expand_dims(state.asnumpy(), axis=0)
          for state in stack_states
      ]
      if is_lstm:
          stack_states_h = stack_states[0::2]
          stack_states_c = stack_states[1::2]
          stack_states = [
              _np.concatenate(stack_states_h, axis=0),
              _np.concatenate(stack_states_c, axis=0),
          ]
      else:
          stack_states = [_np.concatenate(stack_states, axis=0)]
  
      for f, s in zip(fused_states, stack_states):
          assert f.shape == s.shape
          assert_almost_equal(f, s, atol=1e-4, rtol=1e-4)
  
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 1 for calling a global function with example usage: check_rnn_states(fused_states, stack_states, num_layers, bidirectional, is_lstm) without return types