code: def test_t5_model(cfg_key, activation, ctx):
      with ctx:
          cfg = T5Model.get_cfg(cfg_key)
          cfg.defrost()
          cfg.MODEL.vocab_size = 256
          cfg.MODEL.d_model = 128
          cfg.MODEL.d_ff = 512
          cfg.MODEL.num_layers = 2
          cfg.MODEL.num_heads = 4
          cfg.MODEL.activation = activation
          cfg.MODEL.layout = "NT"
          cfg.freeze()
  
          cfg_tn = cfg.clone()
          cfg_tn.defrost()
          cfg_tn.MODEL.layout = "TN"
          cfg_tn.freeze()
  
          # test TN and NT consistency
          t5_model = T5Model.from_cfg(cfg)
          t5_model.initialize()
          t5_model.hybridize()
          t5_model_tn = T5Model.from_cfg(cfg_tn)
          t5_model_tn.share_parameters(t5_model.collect_params())
          t5_model_tn.hybridize()
  
          batch_size = 8
          src_length = 32
          tgt_length = 18
          src_data = np.random.randint(0, 255, (batch_size, src_length))
          src_valid_length = np.random.randint(
              src_length // 2, src_length, (batch_size,)
          )
          tgt_data = np.random.randint(0, 255, (batch_size, tgt_length))
          tgt_valid_length = np.random.randint(
              tgt_length // 4, tgt_length, (batch_size,)
          )
  
          out = t5_model(
              src_data, src_valid_length, tgt_data, tgt_valid_length
          )
          out_tn = t5_model_tn(
              src_data.T, src_valid_length, tgt_data.T, tgt_valid_length
          )
          assert np.allclose(np.swapaxes(out, 0, 1), out_tn, 1e-5, 1e-5)
  
          # test consistency with various target valid length
          for shift in range(1, np.min(tgt_valid_length).item()):
              for partial_out in [
                  t5_model(
                      src_data,
                      src_valid_length,
                      tgt_data[:, :-shift],
                      tgt_valid_length - shift,
                  ),
                  t5_model(
                      src_data,
                      src_valid_length,
                      tgt_data,
                      tgt_valid_length - shift,
                  ),
              ]:
                  for i in range(batch_size):
                      vl = tgt_valid_length[i].item() - shift
                      assert np.allclose(
                          partial_out[i, :vl], out[i, :vl], 1e-5, 1e-5
                      )
  
file_path: gluon-nlp/tests/test_models_t5.py
imports: import pytest
  import tempfile
  import mxnet
  from mxnet import np, npx
  from mxnet.gluon import nn, HybridBlock
  from gluonnlp.models.t5 import T5Model, T5Inference, t5_cfg_reg, list_pretrained_t5, get_pretrained_t5
  from gluonnlp.utils.testing import verify_nmt_model, verify_nmt_inference
label_desc: Use image node_id 2 for calling a global function with example usage: test_t5_model(cfg_key, activation, ctx) without return types