code: def test_layer_attn_probs(self):
      model = TFLongformerModel.from_pretrained(
          "patrickvonplaten/longformer-random-tiny"
      )
      layer = model.longformer.encoder.layer[0].attention.self_attention
      hidden_states = tf.concat(
          [self._get_hidden_states(), self._get_hidden_states() - 0.5],
          axis=0,
      )
      batch_size, seq_length, hidden_size = hidden_states.shape
  
      # create attn mask
      attention_mask_1 = tf.zeros(
          (1, 1, 1, seq_length), dtype=tf.float32
      )
      attention_mask_2 = tf.zeros(
          (1, 1, 1, seq_length), dtype=tf.float32
      )
  
      attention_mask_1 = tf.where(
          tf.range(4)[None, :, None, None] > 1,
          10000.0,
          attention_mask_1,
      )
      attention_mask_1 = tf.where(
          tf.range(4)[None, :, None, None] > 2,
          -10000.0,
          attention_mask_1,
      )
      attention_mask_2 = tf.where(
          tf.range(4)[None, :, None, None] > 0,
          10000.0,
          attention_mask_2,
      )
      attention_mask = tf.concat(
          [attention_mask_1, attention_mask_2], axis=0
      )
  
      is_index_masked = tf.math.less(attention_mask[:, :, 0, 0], 0)
      is_index_global_attn = tf.math.greater(
          attention_mask[:, :, 0, 0], 0
      )
      is_global_attn = tf.math.reduce_any(is_index_global_attn)
  
      layer_head_mask = None
  
      output_hidden_states, local_attentions, global_attentions = layer(
          [
              hidden_states,
              -tf.math.abs(attention_mask),
              layer_head_mask,
              is_index_masked,
              is_index_global_attn,
              is_global_attn,
          ]
      )
  
      self.assertEqual(local_attentions.shape, (2, 4, 2, 8))
      self.assertEqual(global_attentions.shape, (2, 2, 3, 4))
  
      self.assertTrue(
          (local_attentions[0, 2:4, :, :] == 0).numpy().tolist()
      )
      self.assertTrue(
          (local_attentions[1, 1:4, :, :] == 0).numpy().tolist()
      )
  
      #
      # The weight of all tokens with local attention must sum to 1.
      self.assertTrue(
          (
              tf.math.abs(
                  tf.math.reduce_sum(
                      global_attentions[0, :, :2, :], axis=-1
                  )
                  - 1
              )
              < 1e-6
          )
          .numpy()
          .tolist()
      )
      self.assertTrue(
          (
              tf.math.abs(
                  tf.math.reduce_sum(
                      global_attentions[1, :, :1, :], axis=-1
                  )
                  - 1
              )
              < 1e-6
          )
          .numpy()
          .tolist()
      )
  
      tf.debugging.assert_near(
          local_attentions[0, 0, 0, :],
          tf.convert_to_tensor(
              [
                  0.3328,
                  0.0000,
                  0.0000,
                  0.0000,
                  0.0000,
                  0.3355,
                  0.3318,
                  0.0000,
              ],
              dtype=tf.float32,
          ),
          rtol=1e-3,
          atol=1e-4,
      )
  
      tf.debugging.assert_near(
          local_attentions[1, 0, 0, :],
          tf.convert_to_tensor(
              [
                  0.2492,
                  0.2502,
                  0.2502,
                  0.0000,
                  0.0000,
                  0.2505,
                  0.0000,
                  0.0000,
              ],
              dtype=tf.float32,
          ),
          rtol=1e-3,
          atol=1e-4,
      )
  
      # All the global attention weights must sum to 1.
      self.assertTrue(
          (
              tf.math.abs(
                  tf.math.reduce_sum(global_attentions, axis=-1) - 1
              )
              < 1e-6
          )
          .numpy()
          .tolist()
      )
  
      tf.debugging.assert_near(
          global_attentions[0, 0, 1, :],
          tf.convert_to_tensor(
              [0.2500, 0.2500, 0.2500, 0.2500], dtype=tf.float32
          ),
          rtol=1e-3,
          atol=1e-4,
      )
      tf.debugging.assert_near(
          global_attentions[1, 0, 0, :],
          tf.convert_to_tensor(
              [0.2497, 0.2500, 0.2499, 0.2504], dtype=tf.float32
          ),
          rtol=1e-3,
          atol=1e-4,
      )
  
imports: import unittest
  from transformers import is_tf_available
  from transformers.testing_utils import require_sentencepiece, require_tf, require_tokenizers, slow, tooslow
  from ...test_configuration_common import ConfigTester
  from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor, random_attention_mask
label_desc: Use image node_id 8 for calling the TFLongformerModelIntegrationTest obj's underlying member method code with example usage: obj.test_layer_attn_probs() without return types