code: def test_robert_small_config(compute_layout, ctx):
      with ctx:
          cfg = RobertaModel.get_cfg()
          cfg.defrost()
          cfg.MODEL.vocab_size = 1000
          cfg.MODEL.num_layers = 2
          cfg.MODEL.hidden_size = 128
          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()
  
          batch_size = 4
          sequence_length = 16
          num_mask = 3
          inputs = mx.np.random.randint(
              0, 10, (batch_size, sequence_length)
          )
          valid_length = mx.np.random.randint(
              3, sequence_length, (batch_size,)
          )
          masked_positions = mx.np.random.randint(
              0, 3, (batch_size, num_mask)
          )
  
          roberta_model = RobertaModel.from_cfg(cfg)
          roberta_model.initialize()
          roberta_model.hybridize()
          contextual_embeddings, pooled_out = roberta_model(
              inputs, valid_length
          )
          roberta_model_tn = RobertaModel.from_cfg(cfg_tn)
          roberta_model_tn.share_parameters(
              roberta_model.collect_params()
          )
          roberta_model_tn.hybridize()
          contextual_embeddings_tn, pooled_out_tn = roberta_model_tn(
              inputs.T, valid_length
          )
          assert_allclose(
              np.swapaxes(contextual_embeddings_tn.asnumpy(), 0, 1),
              contextual_embeddings.asnumpy(),
              1e-3,
              1e-3,
          )
          assert_allclose(
              pooled_out_tn.asnumpy(), pooled_out.asnumpy(), 1e-3, 1e-3
          )
  
          # Test for RobertaForMLM
          roberta_mlm_model = RobertaForMLM(cfg)
          roberta_mlm_model.initialize()
          roberta_mlm_model.hybridize()
          (
              contextual_embedding,
              pooled_out,
              mlm_score,
          ) = roberta_mlm_model(inputs, valid_length, masked_positions)
          roberta_mlm_model_tn = RobertaForMLM(cfg_tn)
          roberta_mlm_model_tn.share_parameters(
              roberta_mlm_model.collect_params()
          )
          roberta_mlm_model_tn.hybridize()
          (
              contextual_embedding_tn,
              pooled_out_tn,
              mlm_score_tn,
          ) = roberta_mlm_model_tn(
              inputs.T, valid_length.T, masked_positions
          )
          assert_allclose(
              np.swapaxes(contextual_embedding_tn.asnumpy(), 0, 1),
              contextual_embedding.asnumpy(),
              1e-3,
              1e-3,
          )
          assert_allclose(
              pooled_out_tn.asnumpy(), pooled_out.asnumpy(), 1e-3, 1e-3
          )
          assert_allclose(
              mlm_score_tn.asnumpy(), mlm_score.asnumpy(), 1e-3, 1e-3
          )
  
          # Test for fp16
          if ctx.device_type == "gpu":
              verify_backbone_fp16(
                  model_cls=RobertaModel,
                  cfg=cfg,
                  ctx=ctx,
                  inputs=[inputs, valid_length],
              )
  
file_path: gluon-nlp/tests/test_models_roberta.py
imports: import pytest
  import numpy
  import mxnet
  import tempfile
  from numpy.testing import assert_allclose
  from gluonnlp.models.roberta import RobertaModel, RobertaForMLM, list_pretrained_roberta, get_pretrained_roberta
  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_robert_small_config(compute_layout, ctx) without return types