class_docstr: ## Attention with Linear Biases (ALiBi)
  
  We override [Multi-Head Attention](../mha.html).
class_docstr_tok: ["#","#","Attention","with","Linear","Biases","(","ALiBi",")","We","override","[","Multi-Head","Attention","]","(","..","\/mha.html",")","."]
docstr: `query`, `key` and `value` are the tensors that store
  collection of *query*, *key* and *value* vectors.
  They have shape `[seq_len, batch_size, d_model]`.
  
  `mask` has shape `[seq_len, seq_len, batch_size]` and
  `mask[i, j, b]` indicates whether for batch `b`,
  query at position `i` has access to key-value at position `j`.
docstr_tok: ["`","query","`",",","`","key","`","and","`","value","`","are","the","tensors","that","store","collection","of","*","query","*",",","*","key","*","and","*","value","*","vectors",".","They","have","shape","`","[","seq_len",",","batch_size",",","d_model","]","`",".","`","mask","`","has","shape","`","[","seq_len",",","seq_len",",","batch_size","]","`","and","`","mask","[","i",",","j",",","b","]","`","indicates","whether","for","batch","`","b","`",",","query","at","position","`","i","`","has","access","to","key-value","at","position","`","j","`","."]
code: def forward(
      self,
      *,
      query: torch.Tensor,
      key: torch.Tensor,
      value: torch.Tensor,
      mask: Optional[torch.Tensor] = None,
  ):
      """
      `query`, `key` and `value` are the tensors that store
      collection of *query*, *key* and *value* vectors.
      They have shape `[seq_len, batch_size, d_model]`.
  
      `mask` has shape `[seq_len, seq_len, batch_size]` and
      `mask[i, j, b]` indicates whether for batch `b`,
      query at position `i` has access to key-value at position `j`.
      """
  
      # ALiBi only works with causal masks.
      assert mask is not None
      assert mask.shape[0] == mask.shape[1] and mask.shape[2] == 1
  
      # `query`, `key` and `value` have shape `[seq_len, batch_size, d_model]`
      seq_len, batch_size, _ = query.shape
  
      # Add head dimension to mask and check its shape.
      mask = self.prepare_mask(mask, query.shape, key.shape)
  
      # Prepare `query`, `key` and `value` for attention computation.
      # These will then have shape `[seq_len, batch_size, heads, d_k]`.
      query = self.query(query)
      key = self.key(key)
      value = self.value(value)
  
      # Compute attention scores $Q K^\top$.
      # This gives a tensor of shape `[seq_len, seq_len, batch_size, heads]`.
      scores = self.get_scores(query, key)
  
      # Scale scores $\frac{Q K^\top}{\sqrt{d_k}}$
      scores *= self.scale
  
      # Create AliBi biases if it's not cached
      if (
          self.alibi_biases is None
          or self.alibi_biases.shape[1] < seq_len
      ):
          # `mask` has shape `[seq_len, seq_len, 1, 1]`
          self.alibi_biases = get_alibi_biases(
              scores.shape[-1], mask[:, :, 0, 0]
          )
  
      # Add AliBi biases to attention scores.
      # ALiBi biases has shape `[seq_len, seq_len, n_heads]`
      # and `scores` has shape `[seq_len, seq_len, batch_size, n_heads]`
      scores += self.alibi_biases[:seq_len, :seq_len, None, :]
  
      # Apply mask
      scores = scores.masked_fill(mask == 0, float("-inf"))
  
      # $softmax$ attention along the key sequence dimension
      # $\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)$
      attn = self.softmax(scores)
  
      # Apply dropout
      attn = self.dropout(attn)
  
      # Multiply by values
      # $$\underset{seq}{softmax}\Bigg(\frac{Q K^\top}{\sqrt{d_k}}\Bigg)V$$
      x = torch.einsum("ijbh,jbhd->ibhd", attn, value)
  
      # Concatenate multiple heads
      x = x.reshape(seq_len, batch_size, -1)
  
      # Output layer
      return self.output(x)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/transformers/alibi/__init__.py
imports: import math
  from typing import Optional
  import torch
  from torch import nn
  from labml.logger import inspect
  from labml_nn.transformers.mha import MultiHeadAttention
label_desc: Use image node_id 2 for calling the AlibiMultiHeadAttention obj's underlying member method code with example usage: obj.forward() and returns: self