import tensorflow as tf



def get_assignment_map_from_checkpoint(tvars, init_checkpoint, num_of_group=0):
  """Compute the union of the current variables and checkpoint variables."""
  assignment_map = {}
  initialized_variable_names = {}

  name_to_variable = collections.OrderedDict()
  for var in tvars:
    name = var.name
    m = re.match("^(.*):\\d+$", name)
    if m is not None:
      name = m.group(1)
    name_to_variable[name] = var
  init_vars = tf.train.list_variables(init_checkpoint)
  init_vars_name = [name for (name, _) in init_vars]

  if num_of_group > 0:
    assignment_map = []
    for gid in range(num_of_group):
      assignment_map.append(collections.OrderedDict())
  else:
    assignment_map = collections.OrderedDict()

  for name in name_to_variable:
    if name in init_vars_name:
      tvar_name = name
    elif (re.sub(r"/group_\d+/", "/group_0/",
