code: def verify_nmt_inference(
      train_model,
      inference_model,
      batch_size=4,
      src_seq_length=5,
      tgt_seq_length=10,
      atol=1e-4,
      rtol=1e-3,
  ):
      """Verify the correctness of an NMT inference model. Raise error message if it detects
      any problems.
  
      Parameters
      ----------
      train_model
          The training model
      inference_model
          The inference model
      batch_size
          Batch size
      src_seq_length
          Length of the source sequence
      tgt_seq_length
          Length of the target sequence
      atol
          Absolute tolerance
      rtol
          Relative tolerance
  
      """
      if train_model.layout == "NT":
          src_word_sequences = mx.np.random.randint(
              0,
              train_model.src_vocab_size,
              (batch_size, src_seq_length),
          )
          tgt_word_sequences = mx.np.random.randint(
              0,
              train_model.tgt_vocab_size,
              (batch_size, tgt_seq_length),
          )
      else:
          src_word_sequences = mx.np.random.randint(
              0,
              train_model.src_vocab_size,
              (src_seq_length, batch_size),
          )
          tgt_word_sequences = mx.np.random.randint(
              0,
              train_model.tgt_vocab_size,
              (tgt_seq_length, batch_size),
          )
      src_valid_length = mx.np.random.randint(
          1, src_seq_length, (batch_size,)
      )
      min_tgt_seq_length = max(1, tgt_seq_length - 5)
      tgt_valid_length = mx.np.random.randint(
          min_tgt_seq_length, tgt_seq_length, (batch_size,)
      )
      full_out = train_model(
          src_word_sequences,
          src_valid_length,
          tgt_word_sequences,
          tgt_valid_length,
      )
      if train_model.layout == "NT":
          for partial_batch_size in range(1, batch_size + 1):
              step_out_l = []
              states = inference_model.init_states(
                  src_word_sequences[:partial_batch_size, :],
                  src_valid_length[:partial_batch_size],
              )
              assert is_match_states_batch_size(
                  states,
                  inference_model.state_batch_axis,
                  partial_batch_size,
              )
              for i in range(min_tgt_seq_length):
                  step_out, states = inference_model(
                      tgt_word_sequences[:partial_batch_size, i], states
                  )
                  step_out_l.append(step_out)
              partial_out = mx.np.stack(step_out_l, axis=1)
              npt.assert_allclose(
                  full_out[
                      :partial_batch_size, :min_tgt_seq_length
                  ].asnumpy(),
                  partial_out[:partial_batch_size, :].asnumpy(),
                  atol,
                  rtol,
              )
      elif train_model.layout == "TN":
          for partial_batch_size in range(1, batch_size + 1):
              step_out_l = []
              states = inference_model.init_states(
                  src_word_sequences[:, :partial_batch_size],
                  src_valid_length[:partial_batch_size],
              )
              assert is_match_states_batch_size(
                  states,
                  inference_model.state_batch_axis,
                  partial_batch_size,
              )
              for i in range(min_tgt_seq_length):
                  step_out, states = inference_model(
                      tgt_word_sequences[i, :partial_batch_size], states
                  )
                  step_out_l.append(step_out)
              partial_out = mx.np.stack(step_out_l, axis=0)
              npt.assert_allclose(
                  full_out[
                      :min_tgt_seq_length, :partial_batch_size
                  ].asnumpy(),
                  partial_out[:, :partial_batch_size].asnumpy(),
                  atol,
                  rtol,
              )
      else:
          raise NotImplementedError
  
file_path: gluon-nlp/src/gluonnlp/utils/testing.py
imports: import numpy.testing
  import numpy
  import mxnet
  from mxnet.util import use_np
  from .parameter import move_to_ctx
label_desc: Use image node_id 3 for calling a global function with example usage: verify_nmt_inference(train_model, inference_model, batch_size, src_seq_length, tgt_seq_length, atol, rtol) without return types