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"]
code: def _dropout(self, context, x):
      if context.train and self._dropout_rate > 0:
          return mtf.dropout(
              x,
              context.train,
              rate=self._dropout_rate,
              noise_shape=mtf.Shape(
                  context.batch_dims + [context.model.model_dim]
              ),
          )
      else:
          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 14 for calling the UTLayerStack obj's underlying member method code with example usage: obj._dropout(context, x) and returns: mtf, x