class_docstr: A `Trimmer` that allocates a length budget to segments in order.
  
  A `Trimmer` that allocates a length budget to segments in order. It selects
  elements to drop, according to a max sequence length budget, and then applies
  this mask to actually drop the elements. See `generate_mask()` for more
  details.
  
  Example:
  
  >>> a = tf.ragged.constant([['a', 'b', 'c'], [], ['d']])
  >>> b = tf.ragged.constant([['1', '2', '3'], [], ['4', '5', '6', '7']])
  >>> trimmer = tf_text.WaterfallTrimmer(4)
  >>> trimmer.trim([a, b])
  [<tf.RaggedTensor [[b'a', b'b', b'c'], [], [b'd']]>,
   <tf.RaggedTensor [[b'1'], [], [b'4', b'5', b'6']]>]
  
  Here, for the first pair of elements, `['a', 'b', 'c']` and `['1', '2', '3']`,
  the `'2'` and `'3'` are dropped to fit the sequence within the max sequence
  length budget.
class_docstr_tok: ["A","`","Trimmer","`","that","allocates","a","length","budget","to","segments","in","order",".","A","`","Trimmer","`","that","allocates","a","length","budget","to","segments","in","order",".","It","selects","elements","to","drop",",","according","to","a","max","sequence","length","budget",",","and","then","applies","this","mask","to","actually","drop","the","elements",".","See","`","generate_mask","(",")","`","for","more","details",".","Example",":",">",">",">","a","=","tf.ragged.constant","(","[","[","'","a","'",",","'","b","'",",","'","c","'","]",",","[","]",",","[","'d","'","]","]",")",">",">",">","b","=","tf.ragged.constant","(","[","[","'","1","'",",","'","2","'",",","'","3","'","]",",","[","]",",","[","'","4","'",",","'","5","'",",","'","6","'",",","'","7","'","]","]",")",">",">",">","trimmer","=","tf_text.WaterfallTrimmer","(","4",")",">",">",">","trimmer.trim","(","[","a",",","b","]",")","[","<","tf.RaggedTensor","[","[","b","'","a","'",",","b","'","b","'",",","b","'","c","'","]",",","[","]",",","[","b'd","'","]","]",">",",","<","tf.RaggedTensor","[","[","b","'","1","'","]",",","[","]",",","[","b","'","4","'",",","b","'","5","'",",","b","'","6","'","]","]",">","]","Here",",","for","the","first","pair","of","elements",",","`","[","'","a","'",",","'","b","'",",","'","c","'","]","`","and","`","[","'","1","'",",","'","2","'",",","'","3","'","]","`",",","the","`","'","2","'","`","and","`","'","3","'","`","are","dropped","to","fit","the","sequence","within","the","max","sequence","length","budget","."]
docstr: Calculates a truncation mask given a per-batch budget.
  
  Calculate a truncation mask given a budget of the max number of items for
  each or all batch row. The allocation of the budget is done using a
  'waterfall' algorithm. This algorithm allocates quota in a left-to-right
  manner and fill up the buckets until we run out of budget.
  
  For example if the budget of [5] and we have segments of size
  [3, 4, 2], the truncate budget will be allocated as [3, 2, 0].
  
  The budget can be a scalar, in which case the same budget is broadcasted
  and applied to all batch rows. It can also be a 1D `Tensor` of size
  `batch_size`, in which each batch row i will have a budget corresponding to
  `per_batch_quota[i]`.
  
  Example:
  
  >>> a = tf.ragged.constant([['a', 'b', 'c'], [], ['d']])
  >>> b = tf.ragged.constant([['1', '2', '3'], [], ['4', '5', '6', '7']])
  >>> trimmer = tf_text.WaterfallTrimmer(4)
  >>> trimmer.generate_mask([a, b])
  [<tf.RaggedTensor [[True, True, True], [], [True]]>,
   <tf.RaggedTensor [[True, False, False], [], [True, True, True, False]]>]
  
  Args:
    segments: A list of `RaggedTensor` each w/ a shape of [num_batch,
      (num_items)].
  
  Returns:
    a list with len(segments) of `RaggedTensor`s, see superclass for details.
docstr_tok: ["Calculates","a","truncation","mask","given","a","per-batch","budget",".","Calculate","a","truncation","mask","given","a","budget","of","the","max","number","of","items","for","each","or","all","batch","row",".","The","allocation","of","the","budget","is","done","using","a","'waterfall","'","algorithm",".","This","algorithm","allocates","quota","in","a","left-to-right","manner","and","fill","up","the","buckets","until","we","run","out","of","budget",".","For","example","if","the","budget","of","[","5","]","and","we","have","segments","of","size","[","3",",","4",",","2","]",",","the","truncate","budget","will","be","allocated","as","[","3",",","2",",","0","]",".","The","budget","can","be","a","scalar",",","in","which","case","the","same","budget","is","broadcasted","and","applied","to","all","batch","rows",".","It","can","also","be","a","1D","`","Tensor","`","of","size","`","batch_size","`",",","in","which","each","batch","row","i","will","have","a","budget","corresponding","to","`","per_batch_quota","[","i","]","`",".","Example",":",">",">",">","a","=","tf.ragged.constant","(","[","[","'","a","'",",","'","b","'",",","'","c","'","]",",","[","]",",","[","'d","'","]","]",")",">",">",">","b","=","tf.ragged.constant","(","[","[","'","1","'",",","'","2","'",",","'","3","'","]",",","[","]",",","[","'","4","'",",","'","5","'",",","'","6","'",",","'","7","'","]","]",")",">",">",">","trimmer","=","tf_text.WaterfallTrimmer","(","4",")",">",">",">","trimmer.generate_mask","(","[","a",",","b","]",")","[","<","tf.RaggedTensor","[","[","True",",","True",",","True","]",",","[","]",",","[","True","]","]",">",",","<","tf.RaggedTensor","[","[","True",",","False",",","False","]",",","[","]",",","[","True",",","True",",","True",",","False","]","]",">","]","Args",":","segments",":","A","list","of","`","RaggedTensor","`","each","w\/","a","shape","of","[","num_batch",",","(","num_items",")","]",".","Returns",":","a","list","with","len","(","segments",")","of","`","RaggedTensor","`","s",",","see","superclass","for","details","."]
code: def generate_mask(self, segments):
      """Calculates a truncation mask given a per-batch budget.
  
      Calculate a truncation mask given a budget of the max number of items for
      each or all batch row. The allocation of the budget is done using a
      'waterfall' algorithm. This algorithm allocates quota in a left-to-right
      manner and fill up the buckets until we run out of budget.
  
      For example if the budget of [5] and we have segments of size
      [3, 4, 2], the truncate budget will be allocated as [3, 2, 0].
  
      The budget can be a scalar, in which case the same budget is broadcasted
      and applied to all batch rows. It can also be a 1D `Tensor` of size
      `batch_size`, in which each batch row i will have a budget corresponding to
      `per_batch_quota[i]`.
  
      Example:
  
      >>> a = tf.ragged.constant([['a', 'b', 'c'], [], ['d']])
      >>> b = tf.ragged.constant([['1', '2', '3'], [], ['4', '5', '6', '7']])
      >>> trimmer = tf_text.WaterfallTrimmer(4)
      >>> trimmer.generate_mask([a, b])
      [<tf.RaggedTensor [[True, True, True], [], [True]]>,
       <tf.RaggedTensor [[True, False, False], [], [True, True, True, False]]>]
  
      Args:
        segments: A list of `RaggedTensor` each w/ a shape of [num_batch,
          (num_items)].
  
      Returns:
        a list with len(segments) of `RaggedTensor`s, see superclass for details.
      """
      with ops.name_scope("WaterfallTrimmer/generate_mask"):
          segment_row_lengths = [
              _get_row_lengths(s, self._axis) for s in segments
          ]
          segment_row_lengths = array_ops_stack.stack(
              segment_row_lengths, axis=-1
          )
  
          # Broadcast budget to match the rank of segments[0]
          budget = ops.convert_to_tensor(self._max_seq_length)
          for _ in range(segments[0].shape.ndims - budget.shape.ndims):
              budget = array_ops.expand_dims(budget, -1)
  
          # Compute the allocation for each segment using a `waterfall` algorithm
          segment_lengths = math_ops.cast(
              segment_row_lengths, dtypes.int32
          )
          budget = math_ops.cast(budget, dtypes.int32)
          leftover_budget = (
              math_ops.cumsum(
                  -1 * segment_lengths, exclusive=False, axis=-1
              )
              + budget
          )
          leftover_budget = segment_lengths + math_ops.minimum(
              leftover_budget, 0
          )
          results = math_ops.maximum(leftover_budget, 0)
  
          # Translate the results into boolean masks that match the shape of each
          # segment
          results = array_ops_stack.unstack(results, axis=-1)
          item_selectors = [
              item_selector_ops.FirstNItemSelector(i) for i in results
          ]
          return [
              i.get_selectable(s, self._axis)
              for s, i in zip(segments, item_selectors)
          ]
  
imports: import abc
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.framework import tensor
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops.ragged import ragged_array_ops
  from tensorflow.python.ops.ragged import ragged_tensor
  from tensorflow_text.python.ops import item_selector_ops
  from tensorflow.python.framework import load_library
  from tensorflow.python.platform import resource_loader
  from tensorflow.python.ops import while_loop
label_desc: Use image node_id 2 for calling the WaterfallTrimmer obj's underlying member method code with example usage: obj.generate_mask(segments) and returns: unknown