docstr: For comparison:
  1. Modify the pre-training model `__call__` to skip computing metrics and return masked_lm_output like so:
      ```
      ...
      sequence_output, pooled_output = EncoderModel(
      self.config, random_seed=self.random_seed, name="encoder")(
          input_ids, input_mask, type_ids, deterministic=deterministic)
  
      masked_lm_output = nn.Dense(
          self.config.d_emb,
          kernel_init=default_kernel_init,
          name="predictions_dense")(
              sequence_output)
      masked_lm_output = nn.gelu(masked_lm_output)
      masked_lm_output = nn.LayerNorm(
          epsilon=LAYER_NORM_EPSILON, name="predictions_layer_norm")(
              masked_lm_output)
      masked_lm_logits = layers.OutputProjection(
          kernel=self._get_embedding_table(), name="predictions_output")(
              masked_lm_output)
  
      next_sentence_logits = layers.OutputProjection(
          n_out=2, kernel_init=default_kernel_init, name="classification")(
              pooled_output)
  
      return masked_lm_logits
      ...
      ```
  2. Run the following:
      >>> import jax.numpy as jnp
      >>> import sentencepiece as spm
      >>> from flax.training import checkpoints
      >>> from f_net.models import PreTrainingModel
      >>> from f_net.configs.pretraining import get_config, ModelArchitecture
  
      >>> pretrained_params = checkpoints.restore_checkpoint('./f_net/f_net_checkpoint', None) # Location of original checkpoint
      >>> pretrained_config  = get_config()
      >>> pretrained_config.model_arch = ModelArchitecture.F_NET
  
      >>> vocab_filepath = "./f_net/c4_bpe_sentencepiece.model" # Location of the sentence piece model
      >>> tokenizer = spm.SentencePieceProcessor()
      >>> tokenizer.Load(vocab_filepath)
      >>> with pretrained_config.unlocked():
      >>>     pretrained_config.vocab_size = tokenizer.GetPieceSize()
      >>> tokens = jnp.array([[0, 1, 2, 3, 4, 5]])
      >>> type_ids = jnp.zeros_like(tokens, dtype="i4")
      >>> attention_mask = jnp.ones_like(tokens) # Dummy. This gets deleted inside the model.
  
      >>> flax_pretraining_model = PreTrainingModel(pretrained_config)
      >>> pretrained_model_params = freeze(pretrained_params['target'])
      >>> flax_model_outputs = flax_pretraining_model.apply({"params": pretrained_model_params}, tokens, attention_mask, type_ids, None, None, None, None, deterministic=True)
      >>> masked_lm_logits[:, :3, :3]
docstr_tok: ["For","comparison",":","1",".","Modify","the","pre-training","model","`","__call__","`","to","skip","computing","metrics","and","return","masked_lm_output","like","so",":","``","`","...","sequence_output",",","pooled_output","=","EncoderModel","(","self.config",",","random_seed=self.random_seed",",","name=","''","encoder","''",")","(","input_ids",",","input_mask",",","type_ids",",","deterministic=deterministic",")","masked_lm_output","=","nn.Dense","(","self.config.d_emb",",","kernel_init=default_kernel_init",",","name=","''","predictions_dense","''",")","(","sequence_output",")","masked_lm_output","=","nn.gelu","(","masked_lm_output",")","masked_lm_output","=","nn.LayerNorm","(","epsilon=LAYER_NORM_EPSILON",",","name=","''","predictions_layer_norm","''",")","(","masked_lm_output",")","masked_lm_logits","=","layers.OutputProjection","(","kernel=self._get_embedding_table","(",")",",","name=","''","predictions_output","''",")","(","masked_lm_output",")","next_sentence_logits","=","layers.OutputProjection","(","n_out=2",",","kernel_init=default_kernel_init",",","name=","''","classification","''",")","(","pooled_output",")","return","masked_lm_logits","...","``","`","2",".","Run","the","following",":",">",">",">","import","jax.numpy","as","jnp",">",">",">","import","sentencepiece","as","spm",">",">",">","from","flax.training","import","checkpoints",">",">",">","from","f_net.models","import","PreTrainingModel",">",">",">","from","f_net.configs.pretraining","import","get_config",",","ModelArchitecture",">",">",">","pretrained_params","=","checkpoints.restore_checkpoint","(","'.\/f_net\/f_net_checkpoint","'",",","None",")","#","Location","of","original","checkpoint",">",">",">","pretrained_config","=","get_config","(",")",">",">",">","pretrained_config.model_arch","=","ModelArchitecture.F_NET",">",">",">","vocab_filepath","=","``",".\/f_net\/c4_bpe_sentencepiece.model","''","#","Location","of","the","sentence","piece","model",">",">",">","tokenizer","=","spm.SentencePieceProcessor","(",")",">",">",">","tokenizer.Load","(","vocab_filepath",")",">",">",">","with","pretrained_config.unlocked","(",")",":",">",">",">","pretrained_config.vocab_size","=","tokenizer.GetPieceSize","(",")",">",">",">","tokens","=","jnp.array","(","[","[","0",",","1",",","2",",","3",",","4",",","5","]","]",")",">",">",">","type_ids","=","jnp.zeros_like","(","tokens",",","dtype=","''","i4","''",")",">",">",">","attention_mask","=","jnp.ones_like","(","tokens",")","#","Dummy",".","This","gets","deleted","inside","the","model",".",">",">",">","flax_pretraining_model","=","PreTrainingModel","(","pretrained_config",")",">",">",">","pretrained_model_params","=","freeze","(","pretrained_params","[","'target","'","]",")",">",">",">","flax_model_outputs","=","flax_pretraining_model.apply","(","{","``","params","''",":","pretrained_model_params","}",",","tokens",",","attention_mask",",","type_ids",",","None",",","None",",","None",",","None",",","deterministic=True",")",">",">",">","masked_lm_logits","[",":",",",":3",",",":3","]"]
code: def test_inference_for_masked_lm(self):
      """
      For comparison:
      1. Modify the pre-training model `__call__` to skip computing metrics and return masked_lm_output like so:
          ```
          ...
          sequence_output, pooled_output = EncoderModel(
          self.config, random_seed=self.random_seed, name="encoder")(
              input_ids, input_mask, type_ids, deterministic=deterministic)
  
          masked_lm_output = nn.Dense(
              self.config.d_emb,
              kernel_init=default_kernel_init,
              name="predictions_dense")(
                  sequence_output)
          masked_lm_output = nn.gelu(masked_lm_output)
          masked_lm_output = nn.LayerNorm(
              epsilon=LAYER_NORM_EPSILON, name="predictions_layer_norm")(
                  masked_lm_output)
          masked_lm_logits = layers.OutputProjection(
              kernel=self._get_embedding_table(), name="predictions_output")(
                  masked_lm_output)
  
          next_sentence_logits = layers.OutputProjection(
              n_out=2, kernel_init=default_kernel_init, name="classification")(
                  pooled_output)
  
          return masked_lm_logits
          ...
          ```
      2. Run the following:
          >>> import jax.numpy as jnp
          >>> import sentencepiece as spm
          >>> from flax.training import checkpoints
          >>> from f_net.models import PreTrainingModel
          >>> from f_net.configs.pretraining import get_config, ModelArchitecture
  
          >>> pretrained_params = checkpoints.restore_checkpoint('./f_net/f_net_checkpoint', None) # Location of original checkpoint
          >>> pretrained_config  = get_config()
          >>> pretrained_config.model_arch = ModelArchitecture.F_NET
  
          >>> vocab_filepath = "./f_net/c4_bpe_sentencepiece.model" # Location of the sentence piece model
          >>> tokenizer = spm.SentencePieceProcessor()
          >>> tokenizer.Load(vocab_filepath)
          >>> with pretrained_config.unlocked():
          >>>     pretrained_config.vocab_size = tokenizer.GetPieceSize()
          >>> tokens = jnp.array([[0, 1, 2, 3, 4, 5]])
          >>> type_ids = jnp.zeros_like(tokens, dtype="i4")
          >>> attention_mask = jnp.ones_like(tokens) # Dummy. This gets deleted inside the model.
  
          >>> flax_pretraining_model = PreTrainingModel(pretrained_config)
          >>> pretrained_model_params = freeze(pretrained_params['target'])
          >>> flax_model_outputs = flax_pretraining_model.apply({"params": pretrained_model_params}, tokens, attention_mask, type_ids, None, None, None, None, deterministic=True)
          >>> masked_lm_logits[:, :3, :3]
      """
  
      model = FNetForMaskedLM.from_pretrained("google/fnet-base")
      model.to(torch_device)
  
      input_ids = torch.tensor(
          [[0, 1, 2, 3, 4, 5]], device=torch_device
      )
      with torch.no_grad():
          output = model(input_ids)[0]
  
      vocab_size = 32000
  
      expected_shape = torch.Size((1, 6, vocab_size))
      self.assertEqual(output.shape, expected_shape)
  
      expected_slice = torch.tensor(
          [
              [
                  [-1.7819, -7.7384, -7.5002],
                  [-3.4746, -8.5943, -7.7762],
                  [-3.2052, -9.0771, -8.3468],
              ]
          ],
          device=torch_device,
      )
  
      self.assertTrue(
          torch.allclose(output[:, :3, :3], expected_slice, atol=1e-4)
      )
  
imports: import unittest
  from typing import Dict, List, Tuple
  from transformers import FNetConfig, is_torch_available
  from transformers.models.auto import get_values
  from transformers.testing_utils import require_tokenizers, require_torch, slow, torch_device
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
  from ...test_pipeline_mixin import PipelineTesterMixin
label_desc: Use image node_id 1 for calling the FNetModelIntegrationTest obj's underlying member method code with example usage: obj.test_inference_for_masked_lm() without return types