import re


def strip_var_name(var_name):
  """Strips variable name of sub-strings blocking variable name matching.

  Removes sub-strings that should be ignored when matching checkpointed variable
  names to variable names in the training graph, namely:
  - trailing colon + number, e.g. "W:0" --> "W"
  - partitioning info., e.g. "/a/part_12/b" --> "a/b".
  (Note that checkpointed variables do not have partitioning info in their name,
  while model variables do).

  Args:
    var_name: str, variable name.

  Returns:
    stripped variable name.
  """
  # Strip trailing number, e.g. convert "lstm/W_0:0" to "lstm/W_0".
  var_name = re.sub(r':\d+$', '', var_name)
  # Strip partitioning info, e.g. convert "W_0/part_3/Adagrad" to "W_0/Adagrad".
  var_name = re.sub(r'/part_\d+', '', var_name)
  return var_name