code: def get_alibi_biases(n_heads: int, mask: torch.Tensor):
      """
      ## Calculate the attention biases matrix
  
      * `n_heads` is the number of heads in the attention layer
      * `mask` is the attention mask of shape `[seq_len_q, seq_len_k]`
  
      This returns a matrix of shape `[seq_len_q, seq_len_k, n_heads, ]` with ALiBi attention biases.
      """
  
      # Get slopes $m$ for each head
      m = get_slopes(n_heads).to(mask.device)
  
      # Calculate distances $[0, 1, \dots, N]$
      # Here we calculate the distances using the mask.
      #
      # Since it's causal mask we can just use $[0, 1, \dots, N]$ too.
      # `distance = torch.arange(mask.shape[1], dtype=torch.long, device=mask.device)[None, :]`
      distance = mask.cumsum(dim=-1)
  
      # Multiply them pair-wise to get the AliBi bias matrix
      return distance[:, :, None] * m[None, None, :]
  
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 a global function with example usage: get_alibi_biases(n_heads, mask) and returns: unknown