class_docstr: A `Trimmer` that allocates a length budget to segments via round robin.
  
  A `Trimmer` that allocates a length budget to segments using a round robin
  strategy, then drops elements outside of the segment's allocated budget.
  See `generate_mask()` for more details.
class_docstr_tok: ["A","`","Trimmer","`","that","allocates","a","length","budget","to","segments","via","round","robin",".","A","`","Trimmer","`","that","allocates","a","length","budget","to","segments","using","a","round","robin","strategy",",","then","drops","elements","outside","of","the","segment","'s","allocated","budget",".","See","`","generate_mask","(",")","`","for","more","details","."]
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
  'round robin' algorithm. This algorithm allocates quota in each bucket,
  left-to-right repeatedly until all the buckets are filled.
  
  For example if the budget of [5] and we have segments of size
  [3, 4, 2], the truncate budget will be allocated as [2, 2, 1].
  
  Args:
    segments: A list of `RaggedTensor`s each with 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","'round","robin","'","algorithm",".","This","algorithm","allocates","quota","in","each","bucket",",","left-to-right","repeatedly","until","all","the","buckets","are","filled",".","For","example","if","the","budget","of","[","5","]","and","we","have","segments","of","size","[","3",",","4",",","2","]",",","the","truncate","budget","will","be","allocated","as","[","2",",","2",",","1","]",".","Args",":","segments",":","A","list","of","`","RaggedTensor","`","s","each","with","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
      'round robin' algorithm. This algorithm allocates quota in each bucket,
      left-to-right repeatedly until all the buckets are filled.
  
      For example if the budget of [5] and we have segments of size
      [3, 4, 2], the truncate budget will be allocated as [2, 2, 1].
  
      Args:
        segments: A list of `RaggedTensor`s each with a shape of [num_batch,
          (num_items)].
  
      Returns:
        A list with len(segments) of `RaggedTensor`s, see superclass for details.
      """
      with ops.name_scope("RoundRobinTrimmer/generate_mask"):
          segments = list(
              map(
                  ragged_tensor.convert_to_tensor_or_ragged_tensor,
                  segments,
              )
          )
          # The docs state that the segments argument is required to be a list of
          # RaggedTensors of rank 2. However, the python-only op worked with other
          # ranks, so we continue to execute that code for other ranks so as to not
          # break current models.
          rank_2 = segments[0].shape.ndims == 2
          last_axis = self._axis == -1 or self._axis == 1
          if not last_axis or not rank_2:
              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
              )
              segment_row_lengths = math_ops.cast(
                  segment_row_lengths, dtypes.int32
              )
              budget = ops.convert_to_tensor(self._max_seq_length)
              results = _round_robin_allocation(
                  segment_row_lengths, budget
              )
  
              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)
              ]
          else:
              values = list(map(lambda x: x.values, segments))
              row_splits = list(map(lambda x: x.row_splits, segments))
              o_masks = (
                  gen_trimmer_ops.tf_text_round_robin_generate_masks(
                      self._max_seq_length, values, row_splits
                  )
              )
              return [
                  ragged_tensor.RaggedTensor.from_row_splits(m, s)
                  for m, s in zip(o_masks, row_splits)
              ]
  
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 RoundRobinTrimmer obj's underlying member method code with example usage: obj.generate_mask(segments) and returns: unknown, unknown