nodes: 0, module; 1, function_definition; 2, function_name:pool; 3, parameters; 4, block; 5, identifier:data; 6, identifier:batch_size; 7, identifier:key; 8, default_parameter; 9, default_parameter; 10, default_parameter; 11, default_parameter; 12, expression_statement; 13, if_statement; 14, for_statement; 15, identifier:batch_size_fn; 16, lambda; 17, identifier:random_shuffler; 18, None; 19, identifier:shuffle; 20, False; 21, identifier:sort_within_batch; 22, False; 23, comment:"""Sort within buckets, then batch, then shuffle batches.
  
      Partitions data into chunks of size 100*batch_size, sorts examples within
      each chunk using sort_key, then batch these examples and shuffle the
      batches.
      """; 24, comparison_operator:random_shuffler is None; 25, block; 26, identifier:p; 27, call; 28, block; 29, lambda_parameters; 30, identifier:count; 31, identifier:random_shuffler; 32, None; 33, expression_statement; 34, identifier:batch; 35, argument_list; 36, expression_statement; 37, if_statement; 38, identifier:new; 39, identifier:count; 40, identifier:sofar; 41, assignment; 42, identifier:data; 43, binary_operator:batch_size * 100; 44, identifier:batch_size_fn; 45, assignment; 46, identifier:shuffle; 47, block; 48, else_clause; 49, identifier:random_shuffler; 50, attribute; 51, identifier:batch_size; 52, integer:100; 53, identifier:p_batch; 54, conditional_expression:batch(sorted(p, key=key), batch_size, batch_size_fn) \
              if sort_within_batch \
              else batch(p, batch_size, batch_size_fn); 55, for_statement; 56, block; 57, identifier:random; 58, identifier:shuffle; 59, call; 60, line_continuation:\; 61, identifier:sort_within_batch; 62, line_continuation:\; 63, call; 64, identifier:b; 65, call; 66, block; 67, for_statement; 68, identifier:batch; 69, argument_list; 70, identifier:batch; 71, argument_list; 72, identifier:random_shuffler; 73, argument_list; 74, expression_statement; 75, identifier:b; 76, call; 77, block; 78, call; 79, identifier:batch_size; 80, identifier:batch_size_fn; 81, identifier:p; 82, identifier:batch_size; 83, identifier:batch_size_fn; 84, call; 85, yield; 86, identifier:list; 87, argument_list; 88, expression_statement; 89, identifier:sorted; 90, argument_list; 91, identifier:list; 92, argument_list; 93, identifier:b; 94, identifier:p_batch; 95, yield; 96, identifier:p; 97, keyword_argument; 98, identifier:p_batch; 99, identifier:b; 100, identifier:key; 101, identifier:key
code: def pool(data, batch_size, key, batch_size_fn=lambda new, count, sofar: count,
           random_shuffler=None, shuffle=False, sort_within_batch=False):
      """Sort within buckets, then batch, then shuffle batches.
  
      Partitions data into chunks of size 100*batch_size, sorts examples within
      each chunk using sort_key, then batch these examples and shuffle the
      batches.
      """
      if random_shuffler is None:
          random_shuffler = random.shuffle
      for p in batch(data, batch_size * 100, batch_size_fn):
          p_batch = batch(sorted(p, key=key), batch_size, batch_size_fn) \
              if sort_within_batch \
              else batch(p, batch_size, batch_size_fn)
          if shuffle:
              for b in random_shuffler(list(p_batch)):
                  yield b
          else:
              for b in list(p_batch):
                  yield b