class_docstr: Helper to remove padding from a tensor before sending to the experts.
  
  The padding is computed for one reference tensor containing the padding mask
  and then can be applied to any other tensor of shape [dim_origin,...].
  
  Ex:
      input = [
        [tok1, tok2],
        [tok3, tok4],
        [0, 0],
        [0, 0],
        [tok5, tok6],
        [0, 0],
      ]
      output = [
        [tok1, tok2],
        [tok3, tok4],
        [tok5, tok6],
      ]
class_docstr_tok: ["Helper","to","remove","padding","from","a","tensor","before","sending","to","the","experts",".","The","padding","is","computed","for","one","reference","tensor","containing","the","padding","mask","and","then","can","be","applied","to","any","other","tensor","of","shape","[","dim_origin",",","...","]",".","Ex",":","input","=","[","[","tok1",",","tok2","]",",","[","tok3",",","tok4","]",",","[","0",",","0","]",",","[","0",",","0","]",",","[","tok5",",","tok6","]",",","[","0",",","0","]",",","]","output","=","[","[","tok1",",","tok2","]",",","[","tok3",",","tok4","]",",","[","tok5",",","tok6","]",",","]"]
docstr: Compute and store the location of the padding.
  
  Args:
    pad_mask (tf.Tensor): Reference padding tensor of shape
      [batch_size,length] or [dim_origin] (dim_origin=batch_size*length)
      containing non-zeros positive values to indicate padding location.
docstr_tok: ["Compute","and","store","the","location","of","the","padding",".","Args",":","pad_mask","(","tf.Tensor",")",":","Reference","padding","tensor","of","shape","[","batch_size",",","length","]","or","[","dim_origin","]","(","dim_origin=batch_size","*","length",")","containing","non-zeros","positive","values","to","indicate","padding","location","."]
code: def __init__(self, pad_mask):
      """Compute and store the location of the padding.
  
      Args:
        pad_mask (tf.Tensor): Reference padding tensor of shape
          [batch_size,length] or [dim_origin] (dim_origin=batch_size*length)
          containing non-zeros positive values to indicate padding location.
      """
      self.nonpad_ids = None
      self.dim_origin = None
  
      with tf.name_scope("pad_reduce/get_ids"):
          pad_mask = tf.reshape(pad_mask, [-1])  # Flatten the batch
          # nonpad_ids contains coordinates of zeros rows (as pad_mask is
          # float32, checking zero equality is done with |x| < epsilon, with
          # epsilon=1e-9 as standard, here pad_mask only contains positive values
          # so tf.abs would be redundant)
          self.nonpad_ids = tf.to_int32(tf.where(pad_mask < 1e-9))
          self.dim_origin = tf.shape(pad_mask)[:1]
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
label_desc: Use image node_id 1 to create a new PadRemover object from inherited base classes: object with example: obj = PadRemover(pad_mask)