def canonicalize_tensor_name(name):
  """Canonicalizes tensor names.

  For an op that produces only one output, we may be refer to its output tensor
  as either "op_name:0" or simply "op_name". This standardizes all internal
  names as "op_name:0" to simplify the logic.

  Args:
    name: Input name to canonicalize.

  Returns:
    The canonicalized input name.
  """
  if ":" not in name:
    return f"{name}:0"
  return name