class_docstr: ## Attention layer
class_docstr_tok: ["#","#","Attention","layer"]
docstr: #### Calculate the causal mask
  
  * `attn` has shape [batch_size, query_seq_len, key_seq_len, n_heads]
code: def _get_mask(self, attn: torch.Tensor):
      """
      #### Calculate the causal mask
  
      * `attn` has shape [batch_size, query_seq_len, key_seq_len, n_heads]
      """
  
      # Query and key lengths
      nq, nk = attn.shape[1:3]
  
      # Create mask
      if (
          self.causal_mask is None
          or self.causal_mask.shape[0] != nq
          or self.causal_mask.shape[1] != nk
          or self.causal_mask.device != attn.device
      ):
          self.causal_mask = torch.triu(
              attn.new_ones([nq, nk], dtype=torch.bool), 1 + nk - nq
          )
  
      # Return from cache
      return self.causal_mask[None, :, :, None]
  
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 2 for calling the AttentionLayer obj's underlying member method code with example usage: obj._get_mask(attn) and returns: self