code: def _round_robin_allocation(row_lengths, max_seq_length):
      """Allocating quota via round robin algorithm."""
      distribution = array_ops.zeros_like(row_lengths)
      i = constant_op.constant(0)
      batch_size = array_ops.shape(row_lengths)[0]
      num_segments = array_ops.shape(row_lengths)[1]
      quota_used = array_ops.zeros([batch_size], dtypes.int32)
      max_seq_length_bc = max_seq_length + 0 * quota_used
  
      def _cond(i, dist, quota_used):
          del i
          have_quota = quota_used < max_seq_length_bc
          have_space = math_ops.reduce_any(dist < row_lengths, 1)
          return math_ops.reduce_any(
              math_ops.logical_and(have_quota, have_space)
          )
  
      def _body(i, dist, quota_used):
          index = math_ops.mod(i, num_segments)
          updates = array_ops.where(
              dist[..., index] < row_lengths[..., index],
              array_ops.ones_like(dist[..., index]),
              array_ops.zeros_like(dist[..., index]),
          )
          scatter_index = array_ops.tile([index], [batch_size])
          scatter_index = array_ops.expand_dims(scatter_index, -1)
          batch_dim = array_ops.reshape(
              math_ops.range(batch_size), [batch_size, 1]
          )
          scatter_index_2d = array_ops.concat(
              [batch_dim, scatter_index], -1
          )
          new_dist = array_ops.tensor_scatter_add(
              dist, scatter_index_2d, updates
          )
          return i + 1, new_dist, quota_used + updates
  
      _, results, _ = while_loop.while_loop(
          _cond, _body, (i, distribution, quota_used)
      )
      return results
  
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 a global function with example usage: _round_robin_allocation(row_lengths, max_seq_length) and returns: results, math_ops, unknown,  new_dist,  unknown