code: def local_heavy_hitter_mask(
      attn_weights, heavy_budget, no_padding_seq_length=None
  ):
      # attn_weights (BS, head, query, keys)
      dtype_attn_weights = attn_weights.dtype
      seq_length = attn_weights.shape[-1]
      if no_padding_seq_length is None:
          padding_length = 0
      else:
          padding_length = seq_length - no_padding_seq_length
  
      offset = torch.finfo(attn_weights.dtype).min
      tmp_attn = nn.functional.softmax(
          attn_weights, dim=-1, dtype=torch.float32
      ).to(dtype_attn_weights)
  
      accumulated_attention_score = torch.sum(
          tmp_attn[
              :, :, padding_length : heavy_budget + padding_length, :
          ],
          dim=-2,
      )  # (head, keys)
      accumulated_attention_score[
          :, :, heavy_budget + padding_length :
      ] = 0
      if padding_length > 0:
          accumulated_attention_score[:, :, :padding_length] = 0
  
      mask_bottom = torch.zeros_like(attn_weights, dtype=torch.bool)
      mask_bottom[
          :,
          :,
          padding_length : heavy_budget + padding_length,
          padding_length : heavy_budget + padding_length,
      ] = True
  
      for token_index in range(
          heavy_budget + padding_length, seq_length
      ):
          tmp_attn_index = nn.functional.softmax(
              attn_weights[:, :, token_index, :],
              dim=-1,
              dtype=torch.float32,
          ).to(dtype_attn_weights)
          _, tmp_topk_index = accumulated_attention_score.topk(
              k=heavy_budget - 1, dim=-1
          )
          zeros_index = torch.zeros_like(
              tmp_attn_index, dtype=torch.bool
          )
          mask_bottom_index = zeros_index.scatter(
              -1, tmp_topk_index, True
          )  # (head, keys)
          mask_bottom_index[:, :, token_index] = True
  
          mask_bottom[:, :, token_index, :] = mask_bottom_index
          accumulated_attention_score += tmp_attn_index
          accumulated_attention_score = (
              accumulated_attention_score * mask_bottom_index
          )
  
      return mask_bottom
  
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 1 for calling a global function with example usage: local_heavy_hitter_mask(attn_weights, heavy_budget, no_padding_seq_length) and returns: mask_bottom