code: def _attn(
      self,
      query,
      key,
      value,
      attention_mask=None,
      head_mask=None,
  ):
      # compute causal mask from causal mask buffer
      query_length, key_length = query.size(-2), key.size(-2)
      causal_mask = self.bias[
          :, :, key_length - query_length : key_length, :key_length
      ]
  
      # Keep the attention weights computation in fp32 to avoid overflow issues
      query = query.to(torch.float32)
      key = key.to(torch.float32)
  
      attn_weights = torch.matmul(query, key.transpose(-1, -2))
  
      mask_value = torch.finfo(attn_weights.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_weights.dtype
      ).to(attn_weights.device)
      attn_weights = torch.where(causal_mask, attn_weights, mask_value)
  
      attn_weights = attn_weights / self.scale_attn
  
      if attention_mask is not None:
          # Apply the attention mask
          attn_weights = attn_weights + attention_mask
  
      attn_weights = nn.functional.softmax(attn_weights, dim=-1)
      attn_weights = attn_weights.to(value.dtype)
      attn_weights = self.attn_dropout(attn_weights)
  
      # 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 torch
  from typing import Optional, Tuple, Union
  from bigdl.llm.transformers.models.utils import init_kv_cache, extend_kv_cache, apply_rotary_pos_emb, append_kv_cache
  from transformers.utils.import_utils import is_torch_fx_proxy
label_desc: Use image node_id 2 for calling a global function with example usage: _attn(query, key, value, attention_mask, head_mask) and returns: attn_output,  attn_weights