code: def get_slopes(n_heads: int):
      """
      ## Get head-specific slope $m$ for each head
  
      * `n_heads` is the number of heads in the attention layer $n$
  
      The slope for first head is
  
      $$\frac{1}{2^{\frac{8}{n}}} = 2^{-\frac{8}{n}}$$
  
      The slopes for the rest of the heads are in a geometric series with a ratio same as above.
  
      For instance when the number of heads is $8$ the slopes are
      $$\frac{1}{2^1}, \frac{1}{2^2}, \dots, \frac{1}{2^8}$$
      """
  
      # Get the closest power of 2 to `n_heads`.
      # If `n_heads` is not a power of 2, then we first calculate slopes to the closest (smaller) power of 2,
      # and then add the remaining slopes.
      n = 2 ** math.floor(math.log2(n_heads))
      # $2^{-\frac{8}{n}}$
      m_0 = 2.0 ** (-8.0 / n)
      # $2^{-1\frac{8}{n}}, 2^{-2 \frac{8}{n}}, 2^{-3 \frac{8}{n}}, \dots$
      m = torch.pow(m_0, torch.arange(1, 1 + n))
  
      # If `n_heads` is not a power of 2, then we add the remaining slopes.
      # We calculate the remaining slopes for $n * 2$ (avoiding slopes added previously).
      # And pick the slopes upto `n_heads`.
      if n < n_heads:
          # $2^{-\frac{8}{2n}}$
          m_hat_0 = 2.0 ** (-4.0 / n)
          # $2^{-1\frac{8}{2n}}, 2^{-3 \frac{8}{2n}}, 2^{-5 \frac{8}{2n}}, \dots$
          # Note that we take steps by $2$ to avoid slopes added previously.
          m_hat = torch.pow(
              m_hat_0, torch.arange(1, 1 + 2 * (n_heads - n), 2)
          )
          # Concatenate the slopes with the remaining slopes.
          m = torch.cat([m, m_hat])
  
      return m
  
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 1 for calling a global function with example usage: get_slopes(n_heads) and returns: m