class_docstr: ## Attention layer
class_docstr_tok: ["#","#","Attention","layer"]
docstr: :param x: has shape `[batch_size, seq_len, n_hidden]`
code: def forward(self, x: torch.Tensor):
      """
      :param x: has shape `[batch_size, seq_len, n_hidden]`
      """
      # Get query, key and value embeddings (all concatenated).
      # The last dimension size will change from n_hidden -> `3 x n_hidden`
      qkv = self.qkv_lin(x)
  
      # Split into heads by changing the shape to `[batch_size, seq_len, n_heads, 3 * d_k]`
      qkv = qkv.view(*qkv.shape[:-1], self.n_heads, -1)
      # Split into query, key and value each of shape `[batch_size, seq_len, n_heads, 3 * d_k]`
      q, k, v = torch.split(qkv, qkv.shape[-1] // 3, dim=-1)
  
      # If we are caching the states of previous tokens
      if get_cache().get("use_cache", False):
          # Get the state id's. We use to retrieve previous states and store the next states
          prev_state_id, next_state_id = get_cache().get("state_ids")
          # If there's cache
          if prev_state_id is not None:
              # Get the past keys and values. These will have shape `[batch_size, prev_seq_len, n_heads, d_k]`
              k_past, v_past = get_cache().pop(
                  f"attn_kv_{prev_state_id}"
              )
              # Offset of the current embeddings
              offset = k_past.shape[1]
  
              # Add RoPE embeddings
              q = self.rope(q, offset=offset)
              k = self.rope(k, offset=offset)
  
              # Concatenate the past
              k = torch.cat([k_past, k], dim=1)
              v = torch.cat([v_past, v], dim=1)
          else:
              # Add RoPE embeddings
              q = self.rope(q)
              k = self.rope(k)
  
          # Save the current state
          get_cache().push(f"attn_kv_{next_state_id}", (k, v))
      else:
          # No cache - simply add RoPE embeddings
          q = self.rope(q)
          k = self.rope(k)
  
      # Use flash attention
      if (
          self.flash_attention is not None
          and q.shape[1] == k.shape[1]
          and q.shape[-1] <= 128
      ):
          output = self.compute_flash_attention(q, k, v)
      # Otherwise, use normal attention
      else:
          output = self.compute_attention(q, k, v)
  
      # Reshape from `[batch_size, seq_len, n_heads, d_k] to `[batch_size, seq_len, n_hidden]`
      output = output.reshape(*x.shape)
  
      # Final linear layer
      return self.output(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 3 for calling the AttentionLayer obj's underlying member method code with example usage: obj.forward(x) and returns: self