Uniquify operation names: 'einsum', 'einsum_1', 'einsum_2', etc.

PiperOrigin-RevId: 219818155

Author: Mesh TensorFlow Team

mesh_tensorflow/ops.py

@@ -350,6 +350,8 @@ def __init__(self):
     self._tensors = []
     self._trainable_variables = []
     self._all_variables = []
+    # Maps a name used in the graph to the next id to use for that name.
+    self._names_in_use = {}
 
   def __repr__(self):
     return self.to_string
@@ -374,6 +376,37 @@ def all_variables(self):
   def to_string(self):
     return "\n".join([op.to_string for op in self.operations])
 
+  def unique_name(self, name, mark_as_used=True):
+    """Like tf.Graph.unique_name, returns a unique operation name for `name`.
+
+    Args:
+      name: The name for an operation.
+      mark_as_used: whether to mark this name as being used.
+
+    Returns:
+      A string to use as the name for the operation.
+    """
+    scope_name = tf.get_variable_scope().name
+    if scope_name:
+      name = scope_name + "/" + name
+
+    # As in TensorFlow, treat names as case insensitive when deciding whether
+    # they are in use.
+    name_key = name.lower()
+    i = self._names_in_use.get(name_key, 0)
+    if mark_as_used:
+      self._names_in_use[name_key] = i + 1
+    if i > 0:
+      base_name_key = name_key
+      while name_key in self._names_in_use:
+        name_key = "%s_%d" % (base_name_key, i)
+        i += 1
+      if mark_as_used:
+        self._names_in_use[name_key] = 1
+      name = "%s_%d" % (name, i-1)
+
+    return name
+
 
 class Lowering(object):
   """Lowering of a Graph from Mesh-TensorFlow to TensorFlow.
@@ -1143,10 +1176,7 @@ def __init__(self, inputs, mesh=None, name=None):
     self._outputs = []
     self._mesh = mesh
     assert name is not None
-    scope_name = tf.get_variable_scope().name
-    if scope_name:
-      name = scope_name + "/" + name
-    self._name = name
+    self._name = mesh.graph.unique_name(name)
     mesh.graph.operations.append(self)
 
   @property

mesh_tensorflow/ops_test.py

@@ -116,6 +116,29 @@ def testGraph(self):
     self.assertLen(graph.trainable_variables, 1)
     self.assertLen(graph.all_variables, 2)
 
+  def testGraphNames(self):
+    # Standard Usage.
+    graph = mtf.Graph()
+    self.assertEqual(graph.unique_name("a"), "a")
+    self.assertEqual(graph.unique_name("a"), "a_1")
+    self.assertEqual(graph.unique_name("a"), "a_2")
+
+    # Edge cases, the user may choose the name "a_1".
+    graph = mtf.Graph()
+    self.assertEqual(graph.unique_name("a"), "a")
+    self.assertEqual(graph.unique_name("a"), "a_1")
+    self.assertEqual(graph.unique_name("a_1"), "a_1_1")
+
+    graph = mtf.Graph()
+    self.assertEqual(graph.unique_name("a"), "a")
+    self.assertEqual(graph.unique_name("a_1"), "a_1")
+    self.assertEqual(graph.unique_name("a"), "a_2")
+
+    # Case insensitive.
+    graph = mtf.Graph()
+    self.assertEqual(graph.unique_name("a"), "a")
+    self.assertEqual(graph.unique_name("A"), "A_1")
+
   @tf.contrib.eager.run_test_in_graph_and_eager_modes()
   def testLowering(self):
     graph = mtf.Graph()