code: def test_transformer_encoder_decoder(
      pre_norm, num_enc_layers, num_dec_layers
  ):
      batch_size = 8
      src_seq_length = 20
      tgt_seq_length = 15
      units = 32
      enc = TransformerEncoder(
          units=units,
          hidden_size=64,
          num_layers=num_enc_layers,
          num_heads=4,
          dropout=0.0,
          pre_norm=pre_norm,
      )
      dec = TransformerDecoder(
          units=units,
          hidden_size=64,
          num_layers=num_dec_layers,
          num_heads=4,
          dropout=0.0,
          pre_norm=pre_norm,
      )
      enc.hybridize()
      # disabled due to two different signatures calling attention_cell in this test
      # dec.hybridize()
      enc.initialize()
      dec.initialize()
      src_data = mx.np.random.normal(
          0, 1, (batch_size, src_seq_length, units)
      )
      src_valid_length = mx.np.random.randint(
          1, src_seq_length, (batch_size,)
      )
      dst_data = mx.np.random.normal(
          0, 1, (batch_size, tgt_seq_length, units)
      )
      dst_valid_length = mx.np.random.randint(
          5, tgt_seq_length, (batch_size,)
      )
      encoded_mem = enc(src_data, src_valid_length)
      full_decode_out = dec(
          dst_data, dst_valid_length, encoded_mem, src_valid_length
      )
  
      # Test for the TN layout
      enc_tn = TransformerEncoder(
          units=units,
          hidden_size=64,
          num_layers=num_enc_layers,
          num_heads=4,
          dropout=0.0,
          pre_norm=pre_norm,
          layout="TN",
      )
      enc_tn.share_parameters(enc.collect_params())
      dec_tn = TransformerDecoder(
          units=units,
          hidden_size=64,
          num_layers=num_dec_layers,
          num_heads=4,
          dropout=0.0,
          pre_norm=pre_norm,
          layout="TN",
      )
      dec_tn.share_parameters(dec.collect_params())
      enc_tn.hybridize()
      dec_tn.hybridize()
      encoded_mem_tn = enc_tn(
          mx.np.swapaxes(src_data, 0, 1), src_valid_length
      )
      full_decode_out_tn = dec_tn(
          mx.np.swapaxes(dst_data, 0, 1),
          dst_valid_length,
          encoded_mem_tn,
          src_valid_length,
      )
      assert_allclose(
          encoded_mem_tn.asnumpy(),
          mx.np.swapaxes(encoded_mem, 0, 1).asnumpy(),
          1e-5,
          1e-5,
      )
      assert_allclose(
          full_decode_out_tn.asnumpy(),
          mx.np.swapaxes(full_decode_out, 0, 1).asnumpy(),
          1e-5,
          1e-5,
      )
  
      # Test the consistency via shifting the data and the valid_length
      for i in range(1, dst_valid_length.asnumpy().min()):
          for partial_decode_out in [
              dec(
                  dst_data[:, :(-i), :],
                  dst_valid_length - i,
                  encoded_mem,
                  src_valid_length,
              ),
              dec(
                  dst_data,
                  dst_valid_length - i,
                  encoded_mem,
                  src_valid_length,
              ),
          ]:
              for b in range(batch_size):
                  vl = dst_valid_length.asnumpy()[b] - i
                  assert_allclose(
                      partial_decode_out.asnumpy()[b, :vl, :],
                      full_decode_out.asnumpy()[b, :vl, :],
                      1e-5,
                      1e-5,
                  )
      # Test the decoder layer
      self_causal_mask = gen_self_attn_mask(
          dst_data, dst_valid_length, attn_type="causal"
      )
      mem_attn_mask = gen_mem_attn_mask(
          encoded_mem, src_valid_length, dst_data, dst_valid_length
      )
      enc_mem_attn_mask = gen_mem_attn_mask(
          encoded_mem, src_valid_length, dst_data[:, 0:1, :], None
      )
      print(enc_mem_attn_mask)
      h_out = dec.layers[0](
          dst_data, encoded_mem, self_causal_mask, mem_attn_mask
      )
      states = dec.layers[0].init_states(
          batch_size, h_out.ctx, h_out.dtype
      )
      h_out_from_incremental = []
      for i in range(tgt_seq_length):
          ele_h_out, states = dec.layers[0].incremental_decode(
              dst_data[:, i, :],
              states,
              encoded_mem,
              src_valid_length,
              enc_mem_attn_mask,
          )
          h_out_from_incremental.append(ele_h_out)
      h_out_from_incremental = mx.np.stack(
          h_out_from_incremental, axis=1
      )
  
      for i in range(batch_size):
          val_length = dst_valid_length[i].asnumpy()
          assert_allclose(
              h_out_from_incremental[i, :val_length, :].asnumpy(),
              h_out[i, :val_length, :].asnumpy(),
              1e-5,
              1e-5,
          )
      # Test for the full decoder
      states = dec.init_states(batch_size, src_data.ctx, src_data.dtype)
      final_out_from_incremental = []
      for i in range(tgt_seq_length):
          ele_final_out, states = dec.incremental_decode(
              dst_data[:, i, :], states, encoded_mem, src_valid_length
          )
          final_out_from_incremental.append(ele_final_out)
      final_out_from_incremental = mx.np.stack(
          final_out_from_incremental, axis=1
      )
      for i in range(batch_size):
          val_length = dst_valid_length[i].asnumpy()
          assert_allclose(
              final_out_from_incremental[i, :val_length, :].asnumpy(),
              full_decode_out[i, :val_length, :].asnumpy(),
              1e-5,
              1e-5,
          )
  
file_path: gluon-nlp/tests/test_models_transformer.py
imports: import numpy
  import mxnet
  import pytest
  from numpy.testing import assert_allclose
  from gluonnlp.models.transformer import TransformerEncoder, TransformerDecoder, TransformerModel, TransformerInference, transformer_cfg_reg
  from gluonnlp.attention_cell import gen_mem_attn_mask, gen_self_attn_mask
  from gluonnlp.utils.testing import verify_nmt_model, verify_nmt_inference
  from gluonnlp.utils.testing import verify_backbone_fp16
  from gluonnlp.utils.parameter import count_parameters, deduplicate_param_dict
label_desc: Use image node_id 1 for calling a global function with example usage: test_transformer_encoder_decoder(pre_norm, num_enc_layers, num_dec_layers) without return types