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))

        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]
        
dpkd/transformers/tests/models/fnet/test_modeling_fnet.py