import re


def _split_sharded_vars(name_shape_map):
  """Split shareded variables.

  Args:
    name_shape_map: A dict from variable name to variable shape.

  Returns:
    not_sharded: Names of the non-sharded variables.
    sharded: Names of the sharded varibales.
  """
  sharded = []
  not_sharded = []
  for name in name_shape_map:
    if re.match(name, '_[0-9]+$'):
      if re.sub('_[0-9]+$', '_1', name) in name_shape_map:
        sharded.append(name)
      else:
        not_sharded.append(name)
    else:
      not_sharded.append(name)
  return not_sharded, sharded