code: def test_albert_backbone(static_alloc, static_shape, compute_layout):
      batch_size = 3
      cfg = get_test_cfg()
      cfg.defrost()
      cfg.MODEL.compute_layout = compute_layout
      cfg.freeze()
      model = AlbertModel.from_cfg(cfg, use_pooler=True)
      model.initialize()
      model.hybridize(
          static_alloc=static_alloc, static_shape=static_shape
      )
      cfg_tn = cfg.clone()
      cfg_tn.defrost()
      cfg_tn.MODEL.layout = "TN"
      cfg_tn.freeze()
      model_tn = AlbertModel.from_cfg(cfg_tn, use_pooler=True)
      model_tn.share_parameters(model.collect_params())
      model_tn.hybridize(
          static_alloc=static_alloc, static_shape=static_shape
      )
  
      for seq_length in [64, 96]:
          valid_length = mx.np.random.randint(
              seq_length // 2, seq_length, (batch_size,)
          )
          inputs = mx.np.random.randint(
              0, cfg.MODEL.vocab_size, (batch_size, seq_length)
          )
          token_types = mx.np.random.randint(
              0, cfg.MODEL.num_token_types, (batch_size, seq_length)
          )
          contextual_embedding, pooled_out = model(
              inputs, token_types, valid_length
          )
          contextual_embedding_tn, pooled_out_tn = model_tn(
              inputs.T, token_types.T, valid_length
          )
          # Verify layout
          assert_allclose(
              np.swapaxes(contextual_embedding_tn.asnumpy(), 0, 1),
              contextual_embedding.asnumpy(),
              1e-4,
              1e-4,
          )
          assert_allclose(
              pooled_out_tn.asnumpy(), pooled_out.asnumpy(), 1e-4, 1e-4
          )
          assert contextual_embedding.shape == (
              batch_size,
              seq_length,
              cfg.MODEL.units,
          )
          assert pooled_out.shape == (batch_size, cfg.MODEL.units)
          # Ensure the embeddings that exceed valid_length are masked
          contextual_embedding_np = contextual_embedding.asnumpy()
          pooled_out_np = pooled_out.asnumpy()
          for i in range(batch_size):
              ele_valid_length = valid_length[i].asnumpy()
              assert_allclose(
                  contextual_embedding_np[i, ele_valid_length:],
                  np.zeros_like(
                      contextual_embedding_np[i, ele_valid_length:]
                  ),
                  1e-5,
                  1e-5,
              )
          # Ensure that the content are correctly masked
          new_inputs = mx.np.concatenate(
              [inputs, inputs[:, :5]], axis=-1
          )
          new_token_types = mx.np.concatenate(
              [token_types, token_types[:, :5]], axis=-1
          )
          new_contextual_embedding, new_pooled_out = model(
              new_inputs, new_token_types, valid_length
          )
          new_contextual_embedding_np = (
              new_contextual_embedding.asnumpy()
          )
          new_pooled_out_np = new_pooled_out.asnumpy()
          for i in range(batch_size):
              ele_valid_length = valid_length[i].asnumpy()
              assert_allclose(
                  new_contextual_embedding_np[i, :ele_valid_length],
                  contextual_embedding_np[i, :ele_valid_length],
                  1e-5,
                  1e-5,
              )
          assert_allclose(new_pooled_out_np, pooled_out_np, 1e-4, 1e-4)
  
file_path: gluon-nlp/tests/test_models_albert.py
imports: import pytest
  import numpy
  from numpy.testing import assert_allclose
  import mxnet
  import tempfile
  from gluonnlp.models.albert import AlbertModel, AlbertForMLM, AlbertForPretrain, list_pretrained_albert, get_pretrained_albert
label_desc: Use image node_id 2 for calling a global function with example usage: test_albert_backbone(static_alloc, static_shape, compute_layout) without return types