class_docstr: Provides learnable NxN relative attention bias.
  
  Attributes:
    num_heads: Number of heads for which to provide relative attention.
    nd_shape: Shape for which to provided relative attention bias. For instance,
      for images we we would provide a 2D shape. Note that batch and feature
      dimensions should be excluded here.
    initializer: Initializer for the bias.
class_docstr_tok: ["Provides","learnable","NxN","relative","attention","bias",".","Attributes",":","num_heads",":","Number","of","heads","for","which","to","provide","relative","attention",".","nd_shape",":","Shape","for","which","to","provided","relative","attention","bias",".","For","instance",",","for","images","we","we","would","provide","a","2D","shape",".","Note","that","batch","and","feature","dimensions","should","be","excluded","here",".","initializer",":","Initializer","for","the","bias","."]
docstr: Computes attention bias based on relative positions.
  
  Content-based relative position attention bias was used in:
    https://arxiv.org/pdf/1803.02155.
  Non-content-based relative position attention bias was used in:
    https://arxiv.org/abs/1606.01933.
  
  Args:
    length: Length of self-attention window for relative attention.
    num_heads: Number of attention heads.
    name: Name of the parameter to be created.
  
  Returns:
    A `[num_heads, length, length]` tensor with queries.
docstr_tok: ["Computes","attention","bias","based","on","relative","positions",".","Content-based","relative","position","attention","bias","was","used","in",":","https",":","\/\/arxiv.org\/pdf\/1803.02155",".","Non-content-based","relative","position","attention","bias","was","used","in",":","https",":","\/\/arxiv.org\/abs\/1606.01933",".","Args",":","length",":","Length","of","self-attention","window","for","relative","attention",".","num_heads",":","Number","of","attention","heads",".","name",":","Name","of","the","parameter","to","be","created",".","Returns",":","A","`","[","num_heads",",","length",",","length","]","`","tensor","with","queries","."]
code: def relative_attn_bias(self, length, num_heads, name):
      """Computes attention bias based on relative positions.
  
      Content-based relative position attention bias was used in:
        https://arxiv.org/pdf/1803.02155.
      Non-content-based relative position attention bias was used in:
        https://arxiv.org/abs/1606.01933.
  
      Args:
        length: Length of self-attention window for relative attention.
        num_heads: Number of attention heads.
        name: Name of the parameter to be created.
  
      Returns:
        A `[num_heads, length, length]` tensor with queries.
      """
      # Actually we need only 2 * length - 1 relative positions, but we need at
      # least another entry as padding for relative shift of each row to the right
      num_rel_pos = 2 * length
  
      rel_bias = self.param(
          name, self.initializer, (self.num_heads, num_rel_pos)
      )
  
      # Now we have to shift in order to compute relative biases.
      # Example: length = 3
      # Say we want:  [[0, 1, 2], [-1, 0, 1], [-2, -1, 0]]
      # Start: [[-2, -1, 0, 1, 2, 3], [-2, -1, 0, 1, 2, 3], [-2, -1, 0, 1, 2, 3]]
      # We linearize: [-2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3]
      # We slice: [-2, -1, 0, 1, 2, 3, -2, -1, 0, 1, 2, 3, -2, -1, 0]
      # We reshape: [[-2, -1, 0, 1, 2], [3, -2, -1, 0, 1], [2, 3, -2, -1, 0]]
      # We slice: [[0, 1, 2], [-1, 0, 1], [-2, -1, 0]]
      # Tadaaa!
  
      # [heads, length * num_rel_pos]
      rel_bias = jnp.tile(rel_bias, [1, length])
  
      # [heads, length * (num_rel_pos - 1)]
      num_rel_pos -= 1
      rel_bias = rel_bias[..., : length * num_rel_pos]
  
      # [heads, length, num_rel_pos - 1]
      # Now every row is shifted by 1 to the right.
      rel_bias = rel_bias.reshape(num_heads, length, num_rel_pos)
  
      # [heads, length, length]
      # Slice the overlapping elements from start.
      rel_bias = rel_bias[..., num_rel_pos - length :]
  
      return rel_bias
  
imports: import functools
  from typing import Callable, Optional, Sequence, Tuple
  import flax.linen
  import jax
  import jax.numpy
  import numpy
  from scenic.model_lib.layers import nn_layers
label_desc: Use image node_id 2 for calling the RelativeAttentionBias obj's underlying member method code with example usage: obj.relative_attn_bias(length, num_heads, name) and returns: rel_bias