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: Preprocess the input at the beginning of each step.
  
  Args:
    context: mtf context
    x: input tensor
    step: step
  
  Returns:
    preprocessed input.
docstr_tok: ["Preprocess","the","input","at","the","beginning","of","each","step",".","Args",":","context",":","mtf","context","x",":","input","tensor","step",":","step","Returns",":","preprocessed","input","."]
code: def step_preprocess(self, context, x, step):
      """Preprocess the input at the beginning of each step.
  
      Args:
        context: mtf context
        x: input tensor
        step: step
  
      Returns:
        preprocessed input.
  
      """
      original_channel_size = x.shape.dims[-1]
  
      if self.add_step_timing_signal:
          x = self.add_step_timing_signal_func(context, x, step)
      if (
          self.add_position_timing_signal
          or self.add_position_timing_signal
      ) and self.add_or_concat_timing_signal == "concat":
          # linear projection to the original dimension of x
          new_dims = x.shape.dims[:-1] + [original_channel_size]
          x = mtf.layers.dense(
              x,
              variable_dtype=context.variable_dtype,
              new_dims=new_dims,
              activation=None,
              use_bias=False,
          )
          # TODO(yanqiz): implement sru in a separate CL
  
      return x
  
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 6 for calling the UTLayerStack obj's underlying member method code with example usage: obj.step_preprocess(context, x, step) and returns: x