nodes: 0, module; 1, function_definition; 2, function_name:build_collate_fn; 3, parameters; 4, block; 5, default_parameter; 6, default_parameter; 7, default_parameter; 8, expression_statement; 9, function_definition; 10, function_definition; 11, function_definition; 12, if_statement; 13, identifier:batch_first; 14, False; 15, identifier:parallel; 16, True; 17, identifier:sort; 18, False; 19, comment:"""
      Factory for collate_fn functions.
  
      :param batch_first: if True returns batches in (batch, seq) format, if
          False returns in (seq, batch) format
      :param parallel: if True builds batches from parallel corpus (src, tgt)
      :param sort: if True sorts by src sequence length within each batch
      """; 20, function_name:collate_seq; 21, parameters; 22, block; 23, function_name:parallel_collate; 24, parameters; 25, block; 26, function_name:single_collate; 27, parameters; 28, block; 29, identifier:parallel; 30, block; 31, else_clause; 32, identifier:seq; 33, expression_statement; 34, expression_statement; 35, expression_statement; 36, expression_statement; 37, expression_statement; 38, for_statement; 39, if_statement; 40, return_statement; 41, identifier:seqs; 42, expression_statement; 43, expression_statement; 44, if_statement; 45, return_statement; 46, identifier:src_seqs; 47, expression_statement; 48, if_statement; 49, return_statement; 50, return_statement; 51, block; 52, comment:"""
          Builds batches for training or inference.
          Batches are returned as pytorch tensors, with padding.
  
          :param seq: list of sequences
          """; 53, assignment; 54, assignment; 55, assignment; 56, assignment; 57, pattern_list; 58, call; 59, block; 60, identifier:batch_first; 61, block; 62, tuple; 63, comment:"""
          Builds batches from parallel dataset (src, tgt), optionally sorts batch
          by src sequence length.
  
          :param seqs: tuple of (src, tgt) sequences
          """; 64, assignment; 65, identifier:sort; 66, block; 67, call; 68, comment:"""
          Builds batches from text dataset, optionally sorts batch by src
          sequence length.
  
          :param src_seqs: source sequences
          """; 69, identifier:sort; 70, block; 71, else_clause; 72, expression_list; 73, identifier:parallel_collate; 74, return_statement; 75, identifier:lengths; 76, list_comprehension; 77, identifier:batch_length; 78, call; 79, identifier:shape; 80, tuple; 81, identifier:seq_tensor; 82, call; 83, identifier:i; 84, identifier:s; 85, identifier:enumerate; 86, argument_list; 87, expression_statement; 88, expression_statement; 89, expression_statement; 90, identifier:seq_tensor; 91, identifier:lengths; 92, pattern_list; 93, call; 94, expression_statement; 95, expression_statement; 96, identifier:tuple; 97, argument_list; 98, expression_statement; 99, block; 100, call; 101, call; 102, identifier:single_collate; 103, call; 104, for_in_clause; 105, identifier:max; 106, argument_list; 107, identifier:batch_length; 108, call; 109, attribute; 110, argument_list; 111, identifier:seq; 112, assignment; 113, call; 114, assignment; 115, identifier:src_seqs; 116, identifier:tgt_seqs; 117, identifier:zip; 118, argument_list; 119, assignment; 120, assignment; 121, list_comprehension; 122, assignment; 123, expression_statement; 124, identifier:collate_seq; 125, argument_list; 126, identifier:tuple; 127, argument_list; 128, identifier:len; 129, argument_list; 130, identifier:s; 131, identifier:seq; 132, identifier:lengths; 133, identifier:len; 134, argument_list; 135, identifier:torch; 136, identifier:full; 137, identifier:shape; 138, attribute; 139, keyword_argument; 140, identifier:end_seq; 141, subscript; 142, attribute; 143, argument_list; 144, identifier:seq_tensor; 145, call; 146, list_splat; 147, pattern_list; 148, call; 149, identifier:tgt_seqs; 150, list_comprehension; 151, call; 152, for_in_clause; 153, pattern_list; 154, call; 155, assignment; 156, identifier:src_seqs; 157, identifier:indices; 158, identifier:s; 159, identifier:seq; 160, identifier:config; 161, identifier:PAD; 162, identifier:dtype; 163, attribute; 164, identifier:lengths; 165, identifier:i; 166, subscript; 167, identifier:copy_; 168, subscript; 169, attribute; 170, argument_list; 171, identifier:seqs; 172, identifier:indices; 173, identifier:src_seqs; 174, identifier:zip; 175, argument_list; 176, subscript; 177, for_in_clause; 178, identifier:collate_seq; 179, argument_list; 180, identifier:s; 181, list; 182, identifier:indices; 183, identifier:src_seqs; 184, identifier:zip; 185, argument_list; 186, identifier:indices; 187, call; 188, identifier:torch; 189, identifier:int64; 190, identifier:seq_tensor; 191, slice; 192, identifier:i; 193, identifier:s; 194, slice; 195, identifier:seq_tensor; 196, identifier:t; 197, list_splat; 198, identifier:tgt_seqs; 199, identifier:idx; 200, identifier:idx; 201, identifier:indices; 202, identifier:s; 203, identifier:src_seqs; 204, identifier:tgt_seqs; 205, list_splat; 206, identifier:range; 207, argument_list; 208, identifier:end_seq; 209, identifier:end_seq; 210, call; 211, call; 212, call; 213, identifier:sorted; 214, argument_list; 215, identifier:sorted; 216, argument_list; 217, identifier:len; 218, argument_list; 219, call; 220, keyword_argument; 221, keyword_argument; 222, call; 223, keyword_argument; 224, keyword_argument; 225, identifier:src_seqs; 226, identifier:enumerate; 227, argument_list; 228, identifier:key; 229, lambda; 230, identifier:reverse; 231, True; 232, identifier:enumerate; 233, argument_list; 234, identifier:key; 235, lambda; 236, identifier:reverse; 237, True; 238, identifier:src_seqs; 239, lambda_parameters; 240, call; 241, identifier:src_seqs; 242, lambda_parameters; 243, call; 244, identifier:item; 245, identifier:len; 246, argument_list; 247, identifier:item; 248, identifier:len; 249, argument_list; 250, subscript; 251, subscript; 252, identifier:item; 253, integer:1; 254, identifier:item; 255, integer:1
code: def build_collate_fn(batch_first=False, parallel=True, sort=False):
      """
      Factory for collate_fn functions.
  
      :param batch_first: if True returns batches in (batch, seq) format, if
          False returns in (seq, batch) format
      :param parallel: if True builds batches from parallel corpus (src, tgt)
      :param sort: if True sorts by src sequence length within each batch
      """
      def collate_seq(seq):
          """
          Builds batches for training or inference.
          Batches are returned as pytorch tensors, with padding.
  
          :param seq: list of sequences
          """
          lengths = [len(s) for s in seq]
          batch_length = max(lengths)
  
          shape = (batch_length, len(seq))
          seq_tensor = torch.full(shape, config.PAD, dtype=torch.int64)
  
          for i, s in enumerate(seq):
              end_seq = lengths[i]
              seq_tensor[:end_seq, i].copy_(s[:end_seq])
  
          if batch_first:
              seq_tensor = seq_tensor.t()
  
          return (seq_tensor, lengths)
  
      def parallel_collate(seqs):
          """
          Builds batches from parallel dataset (src, tgt), optionally sorts batch
          by src sequence length.
  
          :param seqs: tuple of (src, tgt) sequences
          """
          src_seqs, tgt_seqs = zip(*seqs)
          if sort:
              indices, src_seqs = zip(*sorted(enumerate(src_seqs),
                                              key=lambda item: len(item[1]),
                                              reverse=True))
              tgt_seqs = [tgt_seqs[idx] for idx in indices]
  
          return tuple([collate_seq(s) for s in [src_seqs, tgt_seqs]])
  
      def single_collate(src_seqs):
          """
          Builds batches from text dataset, optionally sorts batch by src
          sequence length.
  
          :param src_seqs: source sequences
          """
          if sort:
              indices, src_seqs = zip(*sorted(enumerate(src_seqs),
                                              key=lambda item: len(item[1]),
                                              reverse=True))
          else:
              indices = range(len(src_seqs))
  
          return collate_seq(src_seqs), tuple(indices)
  
      if parallel:
          return parallel_collate
      else:
          return single_collate