class_docstr: ## Attention layer
class_docstr_tok: ["#","#","Attention","layer"]
code: def compute_attention(
      self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor
  ):
      # Disable auto-casting to fp16 for attention computation
      with autocast(enabled=False):
          if q.dtype == torch.float16:
              # Convert to fp32 if the current dtype is fp16
              attn = torch.einsum(
                  "bihk,bjhk->bijh", q.float(), k.float()
              )
          else:
              # Do not cast for bfloat
              attn = torch.einsum("bihk,bjhk->bijh", q, k)
  
          # Scale attention
          attn = attn * self.scale
  
          # Get causal mask
          mask = self._get_mask(attn)
          # Apply mask
          attn.masked_fill_(mask, self.mask_fill)
  
          # Attention softmax
          attn = self.softmax(attn)
  
      # Get attention weighted values
      output = torch.einsum("bijh,bjhk->bihk", attn.to(v.dtype), v)
  
      return output
  
imports: import copy
  import math
  from typing import Dict, Optional, Set, Callable, Any, Generator, Tuple
  import torch
  from torch import nn
  from torch.cuda.amp import autocast
  from labml import monit, logger
  from labml.logger import Text
  from labml_nn.neox import checkpoint
  from labml_nn.neox.utils.cache import get_cache
label_desc: Use image node_id 5 for calling the AttentionLayer obj's underlying member method code with example usage: obj.compute_attention(q, k, v) and returns: output