class_docstr: A stack of layers for Universal Transformer.
  
  This implementation is largely adapted from t2t universal transformer
  implementation. Reference:
  third_party/py/tensor2tensor/models/research
class_docstr_tok: ["A","stack","of","layers","for","Universal","Transformer",".","This","implementation","is","largely","adapted","from","t2t","universal","transformer","implementation",".","Reference",":","third_party\/py\/tensor2tensor\/models\/research"]
docstr: Add n-dimensional embedding as the position (horizontal) timing signal.
  
  Args:
    context: mtf context
    x: a tensor with shape [batch, length, depth]
    step: step
  
  Returns:
    a Tensor with the same shape as x.
docstr_tok: ["Add","n-dimensional","embedding","as","the","position","(","horizontal",")","timing","signal",".","Args",":","context",":","mtf","context","x",":","a","tensor","with","shape","[","batch",",","length",",","depth","]","step",":","step","Returns",":","a","Tensor","with","the","same","shape","as","x","."]
code: def add_position_timing_signal_func(self, context, x, step):
      """Add n-dimensional embedding as the position (horizontal) timing signal.
  
      Args:
        context: mtf context
        x: a tensor with shape [batch, length, depth]
        step: step
  
      Returns:
        a Tensor with the same shape as x.
  
      """
  
      if not self.position_start_index:
          index = 0
  
      elif self.position_start_index == "random":
          # Shift all positions randomly
          # TODO(dehghani): What would be reasonable for max number of shift?
          index = mtf.random_uniform(
              context.mesh,
              [],
              maxval=x.shape.dims[1].size,
              dtype=tf.int32,
          )
  
      elif self.position_start_index == "step":
          # Shift positions based on the step
          if self.recurrence_type == "act":
              num_steps = self.act_max_steps
          else:
              num_steps = self.num_rec_steps
          index = mtf.cast(
              x.shape.dims[1].size * step / num_steps, dtype=tf.int32
          )
  
      length = context.length_dim
      channels = context.model.model_dim
      signal = self.get_timing_signal_1d(
          context, length, channels, start_index=index
      )
  
      if self.add_or_concat_timing_signal == "add":
          x_with_timing = x + mtf.cast(signal, x.dtype)
      # Unimplemented
      if self.add_or_concat_timing_signal == "concat":
          batch_dim = x.shape.dims[0]
          out_shape = mtf.Shape([batch_dim] + signal.shape.dims[1:])
          signal_tiled = mtf.broadcast(signal, out_shape)
          x_with_timing = mtf.concat(
              (x, signal_tiled),
              concat_dim_name=signal_tiled.dimension_names[-1],
          )
  
      return x_with_timing
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the UTLayerStack obj's underlying member method code with example usage: obj.add_position_timing_signal_func(context, x, step) and returns: x_with_timing