code: def test_headmasking(self):
      if not self.test_head_masking:
          return
      (
          config,
          inputs_dict,
      ) = self.model_tester.prepare_config_and_inputs_for_common()
      config.return_dict = True
  
      def _prepare_layer_head_mask(
          i, attention_heads, num_hidden_layers
      ):
          if i == 0:
              return np.concatenate(
                  [
                      np.zeros(1, dtype=jnp.int32),
                      np.ones(attention_heads - 1, dtype=jnp.int32),
                  ]
              )
          if i == num_hidden_layers - 1:
              return np.concatenate(
                  [
                      np.zeros(attention_heads - 1, dtype=jnp.int32),
                      np.ones(1, dtype=jnp.int32),
                  ]
              )
          return np.ones(attention_heads, dtype=jnp.int32)
  
      for model_class in self.all_model_classes:
          model = model_class(config)
  
          inputs_dict["output_attentions"] = True
          inputs_dict["output_hidden_states"] = False
          inputs = self._prepare_for_class(
              inputs_dict, model_class
          ).copy()
          # Prepare head mask
          inputs["head_mask"] = np.stack(
              [
                  _prepare_layer_head_mask(
                      i,
                      config.num_attention_heads,
                      config.num_hidden_layers,
                  )
                  for i in range(config.num_hidden_layers)
              ]
          )
          outputs = model(**inputs)
  
          def _check_attentions_validity(attentions):
              # Remove NaN
              for t in attentions:
                  # Check we don't have more than 25% nans (arbitrary)
                  self.assertLess(np.isnan(t).sum(), t.size / 4)
              attentions = [
                  np.where(np.isnan(t), 0.0, t) for t in attentions
              ]
  
              self.assertAlmostEqual(
                  attentions[0][..., 0, :, :].sum(), 0.0
              )
              self.assertNotEqual(
                  attentions[0][..., -1, :, :].sum(), 0.0
              )
              if (
                  len(attentions) > 2
              ):  # encoder-decodere models have only 2 layers in each modules
                  self.assertNotEqual(
                      attentions[1][..., 0, :, :].sum(), 0.0
                  )
              self.assertAlmostEqual(
                  attentions[-1][..., -2, :, :].sum(), 0.0
              )
              self.assertNotEqual(
                  attentions[-1][..., -1, :, :].sum(), 0.0
              )
  
          if model.config.is_encoder_decoder:
              raise NotImplementedError(
                  "The test has not been implemented for encoder-decoder models yet."
              )
          else:
              _check_attentions_validity(outputs.attentions)
  
imports: import copy
  import inspect
  import json
  import random
  import tempfile
  import unittest
  from typing import List, Tuple
  import numpy
  import transformers
  from huggingface_hub import HfFolder, delete_repo, set_access_token
  from requests.exceptions import HTTPError
  from transformers import BertConfig, is_flax_available, is_torch_available
  from transformers.models.auto import get_values
  from transformers.testing_utils import TOKEN, USER, CaptureLogger, is_pt_flax_cross_test, is_staging_test, require_flax, torch_device
  from transformers.utils import logging
  from transformers.utils.generic import ModelOutput
label_desc: Use image node_id 26 for calling the FlaxModelTesterMixin obj's underlying member method code with example usage: obj.test_headmasking() and returns: np, np, np