nodes: 0, module; 1, function_definition; 2, function_name:compute_memory_contents_under_schedule; 3, parameters; 4, block; 5, identifier:self; 6, identifier:schedule; 7, expression_statement; 8, expression_statement; 9, expression_statement; 10, expression_statement; 11, for_statement; 12, return_statement; 13, comment:"""The in-memory tensors present when executing each operation in schedule.
  
      Simulates running operations in the order given by a schedule. Keeps track
      of the tensors in memory at every point in time, and outputs a list (one
      entry for each point in time) of all sets of all memory contents (i.e. a
      frozenset of strings) ever seen in this execution.
  
      It is assumed (but not checked) that schedule is a valid topological sort of
      the operations in this graph.
  
      Args:
        schedule: A list of integer ids; the order to run operations in.
  
      Returns:
        a list of frozenset of strings, where the ith entry describes the tensors
        in memory when executing operation i (where schedule[i] is an index into
        get_all_operation_names()).
      """; 14, assignment; 15, assignment; 16, assignment; 17, identifier:operation_id; 18, identifier:schedule; 19, block; 20, identifier:memory_contents_for_each_operation; 21, identifier:out_degree; 22, call; 23, identifier:curr_memory_contents; 24, call; 25, identifier:memory_contents_for_each_operation; 26, list; 27, expression_statement; 28, comment:# Allocate new memory to perform the computation at this node.; 29, for_statement; 30, expression_statement; 31, comment:# Free any tensors which are no longer needed.; 32, for_statement; 33, for_statement; 34, attribute; 35, argument_list; 36, identifier:set; 37, argument_list; 38, assignment; 39, identifier:output_name; 40, call; 41, block; 42, call; 43, identifier:output_name; 44, call; 45, block; 46, identifier:input_name; 47, call; 48, block; 49, identifier:self; 50, identifier:_compute_initial_out_degree; 51, identifier:operation_name; 52, attribute; 53, attribute; 54, argument_list; 55, expression_statement; 56, attribute; 57, argument_list; 58, attribute; 59, argument_list; 60, if_statement; 61, attribute; 62, argument_list; 63, expression_statement; 64, if_statement; 65, subscript; 66, identifier:name; 67, identifier:self; 68, identifier:get_operation_output_names; 69, identifier:operation_name; 70, call; 71, identifier:memory_contents_for_each_operation; 72, identifier:append; 73, call; 74, identifier:self; 75, identifier:get_operation_output_names; 76, identifier:operation_name; 77, comparison_operator:out_degree[output_name] == 0; 78, block; 79, identifier:self; 80, identifier:get_operation_input_names; 81, identifier:operation_name; 82, augmented_assignment; 83, comparison_operator:out_degree[input_name] == 0; 84, block; 85, attribute; 86, identifier:operation_id; 87, attribute; 88, argument_list; 89, identifier:frozenset; 90, argument_list; 91, subscript; 92, integer:0; 93, expression_statement; 94, subscript; 95, integer:1; 96, subscript; 97, integer:0; 98, expression_statement; 99, identifier:self; 100, identifier:_operations; 101, identifier:curr_memory_contents; 102, identifier:add; 103, identifier:output_name; 104, identifier:curr_memory_contents; 105, identifier:out_degree; 106, identifier:output_name; 107, call; 108, identifier:out_degree; 109, identifier:input_name; 110, identifier:out_degree; 111, identifier:input_name; 112, call; 113, attribute; 114, argument_list; 115, attribute; 116, argument_list; 117, identifier:curr_memory_contents; 118, identifier:remove; 119, identifier:output_name; 120, identifier:curr_memory_contents; 121, identifier:remove; 122, identifier:input_name
code: def compute_memory_contents_under_schedule(self, schedule):
      """The in-memory tensors present when executing each operation in schedule.
  
      Simulates running operations in the order given by a schedule. Keeps track
      of the tensors in memory at every point in time, and outputs a list (one
      entry for each point in time) of all sets of all memory contents (i.e. a
      frozenset of strings) ever seen in this execution.
  
      It is assumed (but not checked) that schedule is a valid topological sort of
      the operations in this graph.
  
      Args:
        schedule: A list of integer ids; the order to run operations in.
  
      Returns:
        a list of frozenset of strings, where the ith entry describes the tensors
        in memory when executing operation i (where schedule[i] is an index into
        get_all_operation_names()).
      """
      out_degree = self._compute_initial_out_degree()
  
      curr_memory_contents = set()
      memory_contents_for_each_operation = []
  
      for operation_id in schedule:
        operation_name = self._operations[operation_id].name
        # Allocate new memory to perform the computation at this node.
        for output_name in self.get_operation_output_names(operation_name):
          curr_memory_contents.add(output_name)
        memory_contents_for_each_operation.append(frozenset(curr_memory_contents))
  
        # Free any tensors which are no longer needed.
        for output_name in self.get_operation_output_names(operation_name):
          if out_degree[output_name] == 0:
            curr_memory_contents.remove(output_name)
        for input_name in self.get_operation_input_names(operation_name):
          out_degree[input_name] -= 1
          if out_degree[input_name] == 0:
            curr_memory_contents.remove(input_name)
  
      return memory_contents_for_each_operation