code: def verify_nmt_model(
      model,
      batch_size: int = 4,
      src_seq_length: int = 5,
      tgt_seq_length: int = 10,
      atol: float = 1e-4,
      rtol: float = 1e-3,
  ):
      """Verify the correctness of an NMT model. Raise error message if it detects problems.
  
      Parameters
      ----------
      model
          The machine translation model
      batch_size
          The batch size to test the nmt model
      src_seq_length
          Length of the source sequence
      tgt_seq_length
          Length of the target sequence
      atol
          Absolute tolerance.
      rtol
          Relative tolerance.
  
      """
      src_word_sequence = mx.np.random.randint(
          0, model.src_vocab_size, (batch_size, src_seq_length)
      )
      tgt_word_sequence = mx.np.random.randint(
          0, model.tgt_vocab_size, (batch_size, tgt_seq_length)
      )
      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,)
      )
  
      if model.layout == "NT":
          full_out = model(
              src_word_sequence,
              src_valid_length,
              tgt_word_sequence,
              tgt_valid_length,
          )
      else:
          full_out = model(
              src_word_sequence.T,
              src_valid_length,
              tgt_word_sequence.T,
              tgt_valid_length,
          )
          full_out = mx.np.swapaxes(full_out, 0, 1)
      if full_out.shape != (
          batch_size,
          tgt_seq_length,
          model.tgt_vocab_size,
      ):
          raise AssertionError(
              "The output of NMT model does not match the expected output."
              " Model output shape = {}, Expected (B, T, V) = {}".format(
                  full_out.shape,
                  (batch_size, tgt_seq_length, model.tgt_vocab_size),
              )
          )
      for partial_batch_size in range(1, batch_size + 1):
          for i in range(1, min_tgt_seq_length):
              if model.layout == "NT":
                  partial_out = model(
                      src_word_sequence[:partial_batch_size, :],
                      src_valid_length[:partial_batch_size],
                      tgt_word_sequence[:partial_batch_size, :(-i)],
                      tgt_valid_length[:partial_batch_size]
                      - mx.np.array(i, dtype=tgt_valid_length.dtype),
                  )
              else:
                  partial_out = model(
                      src_word_sequence[:partial_batch_size, :].T,
                      src_valid_length[:partial_batch_size],
                      tgt_word_sequence[:partial_batch_size, :(-i)].T,
                      tgt_valid_length[:partial_batch_size]
                      - mx.np.array(i, dtype=tgt_valid_length.dtype),
                  )
                  partial_out = mx.np.swapaxes(partial_out, 0, 1)
              # Verify that the partial output matches the full output
              for b in range(partial_batch_size):
                  partial_vl = tgt_valid_length.asnumpy()[b] - i
                  npt.assert_allclose(
                      full_out[b, :partial_vl].asnumpy(),
                      partial_out[b, :partial_vl].asnumpy(),
                      atol,
                      rtol,
                  )
  
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 2 for calling a global function with example usage: verify_nmt_model(model, batch_size, src_seq_length, tgt_seq_length, atol, rtol) without return types