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: Call the layer stack.
docstr_tok: ["Call","the","layer","stack","."]
code: def call(self, context, x):
      """Call the layer stack."""
      if isinstance(context.sequence_id, mtf.Tensor):
          # We use this mask to zero out the padding regions at each layer.
          # This "fixes" a bug where extreme values leak from the padding into the
          # non-padding regions.
          # TODO(noam): understand this better and make a more principled fix.
          mask = mtf.cast(
              mtf.not_equal(context.sequence_id, 0),
              context.activation_dtype,
          )
      else:
          mask = None
      x = self._dropout(context, x)
      context.layer_outputs.append(x)
      if self.mix_with_transformer_before_ut:
          for _ in range(self.num_vanilla_transformer_layers):
              x = self.vanilla_transformer_layer(context, x, mask)
      # Call a ACT layer
      if self.recurrence_type == "act":
          x = self.act_layer(context, x, mask)
      elif self.recurrence_type == "basic":
          x = self.ut_basic(context, x, mask)
      elif self.recurrence_type == "highway":
          layer_inputs = (x, x, x)
          x = self.ut_highway(context, layer_inputs, mask)
      if self.mix_with_transformer_after_ut:
          for _ in range(self.num_vanilla_transformer_layers):
              x = self.vanilla_transformer_layer(context, x, mask)
      x = self._layer_norm(context, x, name="final_layer_norm")
      x = self._dropout(context, x)
      if mask:
          x *= mask
      context.layer_outputs.append(x)
      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 13 for calling the UTLayerStack obj's underlying member method code with example usage: obj.call(context, x) and returns: x