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 step (vertical) 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","step","(","vertical",")","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_step_timing_signal_func(self, context, x, step):
      """Add n-dimensional embedding as the step (vertical) 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 self.recurrence_type == "act":
          num_steps = self.act_max_steps
      else:
          num_steps = self.num_rec_steps
      channels = x.shape.dims[-1]
  
      if self.step_timing_signal_type == "learned":
          signal = self.get_layer_timing_signal_learned_1d(
              context, channels, step, num_steps
          )
      elif self.step_timing_signal_type == "sinusoid":
          signal = self.get_layer_timing_signal_sinusoid_1d(
              context, channels, step, num_steps
          )
      if self.add_or_concat_timing_signal == "add":
          x_with_timing = x + mtf.cast(signal, x.dtype)
      elif self.add_or_concat_timing_signal == "concat":
          batch_dim = x.shape.dims[0]
          out_shape = mtf.Shape([batch_dim] + x.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 5 for calling the UTLayerStack obj's underlying member method code with example usage: obj.add_step_timing_signal_func(context, x, step) and returns: x_with_timing