code: def test_albert_for_mlm_model(compute_layout):
      batch_size = 3
      cfg = get_test_cfg()
      cfg.defrost()
      cfg.MODEL.compute_layout = compute_layout
      cfg.freeze()
      albert_mlm_model = AlbertForMLM(backbone_cfg=cfg)
      albert_mlm_model.initialize()
      albert_mlm_model.hybridize()
      cfg_tn = cfg.clone()
      cfg_tn.defrost()
      cfg_tn.MODEL.layout = "TN"
      cfg_tn.freeze()
      albert_mlm_tn_model = AlbertForMLM(backbone_cfg=cfg_tn)
      albert_mlm_tn_model.share_parameters(
          albert_mlm_model.collect_params()
      )
      albert_mlm_tn_model.hybridize()
  
      num_mask = 16
      seq_length = 64
      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)
      )
      valid_length = mx.np.random.randint(
          seq_length // 2, seq_length, (batch_size,)
      )
      masked_positions = mx.np.random.randint(
          0, seq_length // 2, (batch_size, num_mask)
      )
      contextual_embeddings, pooled_out, mlm_score = albert_mlm_model(
          inputs, token_types, valid_length, masked_positions
      )
      (
          contextual_embeddings_tn,
          pooled_out_tn,
          mlm_score_tn,
      ) = albert_mlm_tn_model(
          inputs.T, token_types.T, valid_length, masked_positions
      )
      assert_allclose(
          np.swapaxes(contextual_embeddings_tn.asnumpy(), 0, 1),
          contextual_embeddings.asnumpy(),
          1e-4,
          1e-4,
      )
      assert_allclose(
          pooled_out_tn.asnumpy(), pooled_out.asnumpy(), 1e-4, 1e-4
      )
      assert_allclose(
          mlm_score_tn.asnumpy(), mlm_score.asnumpy(), 1e-4, 1e-4
      )
      assert mlm_score.shape == (
          batch_size,
          num_mask,
          cfg.MODEL.vocab_size,
      )
  
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 3 for calling a global function with example usage: test_albert_for_mlm_model(compute_layout) without return types