code: def test_transformer_xl_for_lm(
      cutoffs, div_val, mem_length, query_length
  ):
      vocab_size = 100
      cfg = TransformerXLForLM.get_cfg()
      cfg.defrost()
      cfg.MODEL.vocab_size = vocab_size
      cfg.MODEL.embed_units = 48
      cfg.MODEL.units = 32
      cfg.MODEL.hidden_size = 64
      cfg.MODEL.num_layers = 2
      cfg.MODEL.cutoffs = cutoffs
      cfg.MODEL.div_val = div_val
      cfg.MODEL.layout = "NT"
      cfg.MODEL.dropout = 0.0
      cfg.MODEL.activation_dropout = 0.0
      cfg.MODEL.attention_dropout = 0.0
      cfg.freeze()
      nt_model = TransformerXLForLM(cfg)
      nt_model.initialize()
  
      tn_cfg = cfg.clone()
      tn_cfg.defrost()
      tn_cfg.MODEL.layout = "TN"
      tn_model = TransformerXLForLM(tn_cfg)
      tn_model.initialize()
      for name, param in tn_model.collect_params().items():
          param.set_data(nt_model.collect_params().get(name).data())
      assert_allclose(
          sum(
              mx.np.linalg.norm(param.data()).asnumpy()
              for param in nt_model.collect_params().values()
          ),
          sum(
              mx.np.linalg.norm(param.data()).asnumpy()
              for param in tn_model.collect_params().values()
          ),
      )
      batch_size = 3
      nt_model.set_mem_length(mem_length)
      tn_model.set_mem_length(mem_length)
  
      ctx = mx.cpu()
  
      data = mx.np.random.randint(
          0,
          vocab_size,
          (batch_size, query_length),
          ctx=ctx,
          dtype=np.int32,
      )
      target = mx.np.random.randint(
          0,
          vocab_size,
          (batch_size, query_length),
          ctx=ctx,
          dtype=np.int32,
      )
  
      # Check consistency of layout
      nt_mem_l = nt_model.init_states(batch_size, ctx=ctx)
      for _ in range(8):
          with mx.autograd.record():
              nt_logits, nt_mem_l = nt_model(data, target, nt_mem_l)
              loss = nt_logits.sum()
              loss.backward()
      tn_mem_l = tn_model.init_states(batch_size, ctx=ctx)
      for _ in range(8):
          with mx.autograd.record():
              tn_logits, tn_mem_l = tn_model(data.T, target.T, tn_mem_l)
              loss = tn_logits.sum()
              loss.backward()
      assert_allclose(
          tn_logits.T.asnumpy(), nt_logits.asnumpy(), 1e-5, 1e-5
      )
      for name, tn_param in tn_model.collect_params().items():
          nt_param = nt_model.collect_params().get(name)
          if nt_param.grad_req != "null":
              assert_allclose(
                  nt_param.grad().asnumpy(),
                  tn_param.grad().asnumpy(),
                  1e-4,
                  1e-4,
              )
  
      # Check step_forward consistency
      mem_l = nt_model.init_states(batch_size, ctx=ctx)
      sel_logits, new_mem_l = nt_model(data, target, mem_l)
      ele_sel_logits_l = []
      step_new_mem_l = mem_l
      for i in range(query_length):
          step_logits, step_new_mem_l = nt_model.step_forward(
              data[:, i], step_new_mem_l
          )
          ele_sel_logits_l.append(
              step_logits[mx.np.arange(batch_size), target[:, i]]
          )
      sel_logits_from_step = mx.np.stack(ele_sel_logits_l, axis=-1)
      assert_allclose(
          sel_logits_from_step.asnumpy(),
          sel_logits.asnumpy(),
          1e-4,
          1e-4,
      )
      for lhs, rhs in zip(step_new_mem_l, new_mem_l):
          assert_allclose(lhs.asnumpy(), rhs.asnumpy(), 1e-4, 1e-4)
  
imports: import pytest
  from gluonnlp.models.transformer_xl import TransformerXLForLM
  import mxnet
  import numpy
  from numpy.testing import assert_allclose
  from gluonnlp.utils.parameter import grad_global_norm
label_desc: Use image node_id 1 for calling a global function with example usage: test_transformer_xl_for_lm(cutoffs, div_val, mem_length, query_length) without return types