code: def drop_2d(targets, mode, dropout_p):
      """Dropout in 2D."""
      if dropout_p > 0 and mode == tf_estimator.ModeKeys.TRAIN:
          (
              batch_size,
              targets_length,
              hidden_size,
          ) = common_layers.shape_list(targets)
          mask_prob = tf.random_uniform(
              shape=(batch_size, targets_length), minval=0.0, maxval=1.0
          )
          mask_prob = tf.tile(
              mask_prob[..., tf.newaxis], [1, 1, hidden_size]
          )
          scale = 1 / (1 - dropout_p)
          targets_noisy = tf.where(
              mask_prob > dropout_p,
              targets * scale,
              tf.zeros_like(targets),
          )
          return targets_noisy
      return targets
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from six.moves import range
  from tensor2tensor.layers import common_attention
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers import transformer_glow_layers_ops
  from tensor2tensor.models.transformer import transformer_decoder_layer
  from tensor2tensor.models.transformer import transformer_encoder
  from tensor2tensor.models.transformer import transformer_prepare_encoder
  from tensor2tensor.utils import learning_rate
  from tensor2tensor.utils import mlperf_log
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
label_desc: Use image node_id 7 for calling a global function with example usage: drop_2d(targets, mode, dropout_p) and returns: targets, targets_noisy