code: def _dropout(x, context, dropout_rate):
      if context.train and dropout_rate > 0:
          return mtf.dropout(
              x,
              context.train,
              rate=dropout_rate,
              noise_shape=mtf.Shape(
                  context.batch_dims + x.shape.dims[-1:]
              ),
          )
      else:
          return x
  
imports: import gin
  import mesh_tensorflow
  from mesh_tensorflow.transformer import transformer
  from mesh_tensorflow.transformer import transformer_layers
  import tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: _dropout(x, context, dropout_rate) and returns: mtf, x