code: def _attn(
      self, query, key, value, attention_mask=None, head_mask=None
  ):
      # q, k, v: [bs, num_attention_heads, seq_len, attn_head_size]
      # compute causal mask from causal mask buffer
      (
          batch_size,
          num_attention_heads,
          query_length,
          attn_head_size,
      ) = query.size()
      key_length = key.size(-2)
  
      causal_mask = self.bias[
          :, :, key_length - query_length : key_length, :key_length
      ]
  
      query = query.view(
          batch_size * num_attention_heads, query_length, attn_head_size
      )
      key = key.view(
          batch_size * num_attention_heads, key_length, attn_head_size
      )
      attn_scores = torch.zeros(
          batch_size * num_attention_heads,
          query_length,
          key_length,
          dtype=query.dtype,
          device=key.device,
      )
      attn_scores = torch.baddbmm(
          attn_scores,
          query,
          key.transpose(1, 2),
          beta=1.0,
          alpha=(
              torch.tensor(
                  1.0,
                  dtype=self.norm_factor.dtype,
                  device=self.norm_factor.device,
              )
              / self.norm_factor
          ),
      )
      attn_scores = attn_scores.view(
          batch_size, num_attention_heads, query_length, key_length
      )
  
      mask_value = torch.finfo(attn_scores.dtype).min
      # Need to be a tensor, otherwise we get error: `RuntimeError: expected scalar type float but found double`.
      # Need to be on the same device, otherwise `RuntimeError: ..., x and y to be on the same device`
      mask_value = torch.tensor(mask_value, dtype=attn_scores.dtype).to(
          attn_scores.device
      )
      attn_scores = torch.where(causal_mask, attn_scores, mask_value)
  
      if attention_mask is not None:
          # Apply the attention mask
          attn_scores = attn_scores + attention_mask
  
      if self.attention_masks_next is not None:
          attn_scores = (
              attn_scores * self.attention_masks_next
              + (1 - self.attention_masks_next)
              * torch.finfo(attn_scores.dtype).min
          )
  
      attn_weights = nn.functional.softmax(attn_scores, dim=-1)
      attn_weights = attn_weights.to(value.dtype)
  
      # attn_weights (BS, heads, q-tokens, k-tokens) 16, 15, 15 // 16, 1, 16
      current_scores_sum = attn_weights.sum(0).sum(
          1
      )  # (heads, k-tokens)
      # offset = attn_weights.gt(0).sum(0).sum(1)
  
      # Accumulate attention scores
      if not self.previous_scores == None:
          current_scores_sum[
              :, :-1
          ] += self.previous_scores  # (Enlarge Sequence)
      else:
          self.heavy_budget = int(
              self.heavy_budget_ratio * current_scores_sum.shape[-1]
          )
          self.recent_budget = int(
              self.recent_budget_ratio * current_scores_sum.shape[-1]
          )
          self.cache_budget = self.heavy_budget + self.recent_budget
          self.cache_budget_records.append(self.cache_budget)
          self.input_length.append(attn_weights.shape[-1])
  
      dtype_attn_weights = attn_weights.dtype
      attn_weights_devices = attn_weights.device
      assert attn_weights.shape[0] == 1
      self.previous_scores = current_scores_sum  # (heads, k-tokens)
      attn_mask = (
          torch.ones(
              current_scores_sum.shape[0],
              current_scores_sum.shape[1] + 1,
          )
          .to(dtype_attn_weights)
          .to(attn_weights_devices)
      )
  
      attn_tokens_all = self.previous_scores.shape[-1]
      if attn_tokens_all > self.cache_budget:
          # activate most recent k-cache
          if not self.recent_budget == 0:
              attn_mask[:, : -self.recent_budget] = 0
              selected_set = self.previous_scores[
                  :, : -self.recent_budget
              ]
          else:
              # activate historical best self.cache_budget - self.recent_budget tokens.
              # self.previous_scores # (k-Cache - 1)
              selected_set = self.previous_scores
  
          if not self.heavy_budget == 0:
              _, keep_topk = selected_set.topk(
                  k=self.heavy_budget, dim=-1, largest=True
              )
              attn_mask = attn_mask.scatter(-1, keep_topk, 1)
  
      self.attention_masks_next = attn_mask.unsqueeze(0).unsqueeze(2)
      score_mask = attn_mask[:, :-1]
      score_mask[:, -self.recent_budget :] = 1
      self.previous_scores = self.previous_scores * score_mask
  
      # Mask heads if we want to
      if head_mask is not None:
          attn_weights = attn_weights * head_mask
  
      attn_output = torch.matmul(attn_weights, value)
      return attn_output, attn_weights
  
imports: import os
  import pdb
  import copy
  import math
  import numpy
  from dataclasses import dataclass
  from typing import Optional, Tuple, Union
  import torch
  from torch import nn
  import torch.utils.checkpoint
  import torch.nn.functional
  from torch.cuda.amp import autocast
  from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss
  from transformers.models.gpt_neox.modeling_gpt_neox import RotaryEmbedding, GPTNeoXAttention, apply_rotary_pos_emb
label_desc: Use image node_id 6 for calling the GPTNeoXAttention_Mask obj's underlying member method code with example usage: obj._attn(query, key, value, attention_mask, head_mask) and returns: attn_output,  attn_weights