Adds positional embedding.

  Args:
    x: Tensor with shape [batch, length, depth].
    max_length: int representing static maximum size of any dimension.
    name: str representing name of the embedding tf.Variable.
    positions: Tensor with shape [batch, length].

  Returns:
    Tensor of same shape as x.

def add_positional_embedding(x, max_length, name=None, positions=None):
  """Adds positional embedding.

  Args:
    x: Tensor with shape [batch, length, depth].
    max_length: int representing static maximum size of any dimension.
    name: str representing name of the embedding tf.Variable.
    positions: Tensor with shape [batch, length].

  Returns:
    Tensor of same shape as x.
  """
  with tf.name_scope("add_positional_embedding"):
    _, length, depth = common_layers.shape_list(x)
    var = tf.cast(tf.get_variable(name, [max_length, depth]), x.dtype)
    if positions is None:
      pad_length = tf.maximum(0, length - max_length)
      sliced = tf.cond(
          tf.less(length, max_length),
          lambda: tf.slice(var, [0, 0], [length, -1]),
          lambda: tf.pad(var, [[0, pad_length], [0, 0]]))
      return x + tf.expand_dims(sliced, 0)
    else:
      return x + tf.gather(var, tf.to_int32(positions))