code: def _relative_position_bucket(
      relative_position,
      bidirectional=True,
      num_buckets=32,
      max_distance=128,
  ):
      """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)
      """
      ret = 0
      n = -relative_position
      if bidirectional:
          num_buckets //= 2
          ret += mtf.to_int32(mtf.less(n, 0)) * num_buckets
          n = mtf.abs(n)
      else:
          n = mtf.maximum(n, 0)
      # now n is in the range [0, inf)
      max_exact = num_buckets // 2
      is_small = mtf.less(n, max_exact)
      val_if_large = max_exact + mtf.to_int32(
          mtf.log(mtf.to_float(n) / max_exact)
          / math.log(max_distance / max_exact)
          * (num_buckets - max_exact)
      )
      val_if_large = mtf.minimum(val_if_large, num_buckets - 1)
      ret += mtf.where(is_small, n, val_if_large)
      return ret
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import copy
  import json
  import math
  import re
  import mesh_tensorflow
  import mesh_tensorflow.transformer.moe
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling a global function with example usage: _relative_position_bucket(relative_position, bidirectional, num_buckets, max_distance) and returns: ret