docstr: Adapted from Mesh Tensorflow:
  https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
  
  Translate relative position to a bucket number for relative attention. The relative position is defined as
  memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
  position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
  small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
  positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
  This should allow for more graceful generalization to longer sequences than the model has been trained on
  
  Args:
      relative_position: an int32 Tensor
      bidirectional: a boolean - whether the attention is bidirectional
      num_buckets: an integer
      max_distance: an integer
  
  Returns:
      a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
docstr_tok: ["Adapted","from","Mesh","Tensorflow",":","https",":","\/\/github.com\/tensorflow\/mesh\/blob\/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d\/mesh_tensorflow\/transformer\/transformer_layers.py","#","L593","Translate","relative","position","to","a","bucket","number","for","relative","attention",".","The","relative","position","is","defined","as","memory_position","-","query_position",",","i.e",".","the","distance","in","tokens","from","the","attending","position","to","the","attended-to","position",".","If","bidirectional=False",",","then","positive","relative","positions","are","invalid",".","We","use","smaller","buckets","for","small","absolute","relative_position","and","larger","buckets","for","larger","absolute","relative_positions",".","All","relative","positions",">","=max_distance","map","to","the","same","bucket",".","All","relative","positions","<","=-max_distance","map","to","the","same","bucket",".","This","should","allow","for","more","graceful","generalization","to","longer","sequences","than","the","model","has","been","trained","on","Args",":","relative_position",":","an","int32","Tensor","bidirectional",":","a","boolean","-","whether","the","attention","is","bidirectional","num_buckets",":","an","integer","max_distance",":","an","integer","Returns",":","a","Tensor","with","the","same","shape","as","relative_position",",","containing","int32","values","in","the","range","[","0",",","num_buckets",")"]
code: def _relative_position_bucket(
      relative_position,
      bidirectional=True,
      num_buckets=32,
      max_distance=128,
  ):
      """
      Adapted from Mesh Tensorflow:
      https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593
  
      Translate relative position to a bucket number for relative attention. The relative position is defined as
      memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to
      position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for
      small absolute relative_position and larger buckets for larger absolute relative_positions. All relative
      positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.
      This should allow for more graceful generalization to longer sequences than the model has been trained on
  
      Args:
          relative_position: an int32 Tensor
          bidirectional: a boolean - whether the attention is bidirectional
          num_buckets: an integer
          max_distance: an integer
  
      Returns:
          a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)
      """
      relative_buckets = 0
      if bidirectional:
          num_buckets //= 2
          relative_buckets += (relative_position > 0).to(
              torch.long
          ) * num_buckets
          relative_position = torch.abs(relative_position)
      else:
          relative_position = -torch.min(
              relative_position, torch.zeros_like(relative_position)
          )
      # now relative_position is in the range [0, inf)
  
      # half of the buckets are for exact increments in positions
      max_exact = num_buckets // 2
      is_small = relative_position < max_exact
  
      # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance
      relative_position_if_large = max_exact + (
          torch.log(relative_position.float() / max_exact)
          / math.log(max_distance / max_exact)
          * (num_buckets - max_exact)
      ).to(torch.long)
      relative_position_if_large = torch.min(
          relative_position_if_large,
          torch.full_like(relative_position_if_large, num_buckets - 1),
      )
  
      relative_buckets += torch.where(
          is_small, relative_position, relative_position_if_large
      )
      return relative_buckets
  
imports: import copy
  import math
  import warnings
  from typing import Optional, Tuple, Union
  from loguru import logger
  import torch
  import torch.distributed
  from torch import nn
  from torch.nn import CrossEntropyLoss
  from transformers.activations import ACT2FN
  from transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput
  from transformers.modeling_utils import PreTrainedModel
  from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS
  from transformers.utils import is_torch_fx_proxy
  from transformers import T5Config
  from text_generation_server.utils.layers import TensorParallelColumnLinear, TensorParallelEmbedding, TensorParallelRowLinear, TensorParallelHead
label_desc: Use image node_id 2 for calling the T5Attention obj's underlying member method code with example usage: obj._relative_position_bucket(relative_position, bidirectional, num_buckets, max_distance) and returns: relative_buckets