Rewrite the current graph to combine "Assign" operations.

    Combine similar Assign operations into grouped Assign operations.
    This is useful when using the rewrite_stack_variables() optimization,
    since variables can only be stacked if they are present in the same set
    of Assign operations.

    This function takes a list of Assign operations and returns a possibly
    shorter list of Assign operations.  The input Assignment operations
    are removed from the graph and become invalid.

    Args:
      assignments: a list of Assign objects
    Returns:
      a list of Assign objects

def combine_assignments(self, assignments):
    """Rewrite the current graph to combine "Assign" operations.

    Combine similar Assign operations into grouped Assign operations.
    This is useful when using the rewrite_stack_variables() optimization,
    since variables can only be stacked if they are present in the same set
    of Assign operations.

    This function takes a list of Assign operations and returns a possibly
    shorter list of Assign operations.  The input Assignment operations
    are removed from the graph and become invalid.

    Args:
      assignments: a list of Assign objects
    Returns:
      a list of Assign objects
    """
    group_by_fn = collections.defaultdict(list)
    for a in assignments:
      if not isinstance(a, Assign):
        raise ValueError("ops should be instances of mtf.Assign")
      group_by_fn[a.assign_fn].append(a)
    assignments_set = set(assignments)
    self._operations = [
        op for op in self._operations if op not in assignments_set]
    ret = []
    for fn, ops in six.iteritems(group_by_fn):
      variables = []
      values = []
      for a in ops:
        variables.extend(a.variables)
        values.extend(a.inputs)
      ret.append(Assign(variables, values, fn))
    return ret