code: def test_gpt2_small_config(compute_layout, ctx):
      cfg = GPT2Model.get_cfg()
      cfg.defrost()
      cfg.MODEL.vocab_size = 1000
      cfg.MODEL.units = 128
      cfg.MODEL.num_layers = 2
      cfg.MODEL.num_heads = 2
      cfg.MODEL.compute_layout = compute_layout
      cfg.freeze()
  
      # Generate TN layout
      cfg_tn = cfg.clone()
      cfg_tn.defrost()
      cfg_tn.MODEL.layout = "TN"
      cfg_tn.freeze()
  
      with ctx:
          batch_size = 4
          sequence_length = 16
          inputs = mx.np.random.randint(
              0, 1000, (batch_size, sequence_length), ctx=ctx
          )
  
          gpt2_model = GPT2Model.from_cfg(cfg)
          gpt2_model.initialize(ctx=ctx)
          gpt2_model.hybridize()
          hiddens, _ = gpt2_model(
              inputs, gpt2_model.init_states(batch_size, ctx)
          )
  
          gpt2_model_tn = GPT2Model.from_cfg(cfg_tn)
          gpt2_model_tn.share_parameters(gpt2_model.collect_params())
          gpt2_model_tn.hybridize()
          hiddens_tn, _ = gpt2_model_tn(
              inputs.T, gpt2_model_tn.init_states(batch_size, ctx)
          )
          assert_allclose(
              np.swapaxes(hiddens_tn.asnumpy(), 0, 1),
              hiddens.asnumpy(),
              1e-4,
              1e-4,
          )
  
          # Test for GPT2ForLM
          gpt2_lm_model = GPT2ForLM(cfg)
          gpt2_lm_model.initialize(ctx=ctx)
          gpt2_lm_model.hybridize()
          logits, states = gpt2_lm_model(
              inputs, gpt2_lm_model.init_states(batch_size, ctx)
          )
          gpt2_lm_model_tn = GPT2ForLM(cfg_tn)
          gpt2_lm_model_tn.share_parameters(
              gpt2_lm_model.collect_params()
          )
          gpt2_lm_model_tn.hybridize()
          logits_tn, states_tn = gpt2_lm_model_tn(
              inputs.T, gpt2_lm_model_tn.init_states(batch_size, ctx)
          )
          assert_allclose(
              np.swapaxes(logits_tn.asnumpy(), 0, 1),
              logits.asnumpy(),
              1e-4,
              1e-4,
          )
          assert_allclose(
              np.swapaxes(states_tn.asnumpy(), 2, 3),
              states.asnumpy(),
              1e-4,
              1e-4,
          )
  
          # Verify Float16
          if ctx.device_type == "gpu":
              verify_backbone_fp16(
                  model_cls=GPT2Model,
                  cfg=cfg,
                  ctx=ctx,
                  inputs=[
                      inputs,
                      gpt2_model.init_states(batch_size, ctx),
                  ],
                  check_amp=False,
              )
              pytest.skip(
                  "GPT-2 test has been turned off. "
                  "Issue: https://github.com/apache/incubator-mxnet/issues/19463"
              )
  
file_path: gluon-nlp/tests/test_models_gpt2.py
imports: import pytest
  import numpy
  import mxnet
  import tempfile
  from numpy.testing import assert_allclose
  from gluonnlp.models.gpt2 import GPT2Model, GPT2ForLM, list_pretrained_gpt2, get_pretrained_gpt2
  from gluonnlp.loss import LabelSmoothCrossEntropyLoss
  from gluonnlp.utils.testing import verify_backbone_fp16
label_desc: Use image node_id 2 for calling a global function with example usage: test_gpt2_small_config(compute_layout, ctx) without return types