The first application of a template made with tf.make_template(create_scope_now_=True) will now create its operations in the same name scope as its variables. This means that all variables and ops created by a template that is only applied once end up in the same scope.

Change: 148922089

2 files changed, 52 insertions, 2 deletions

diff --git a/tensorflow/python/kernel_tests/template_test.py b/tensorflow/python/kernel_tests/template_test.py
index 84f46c6278..be2d6a566a 100644
--- a/tensorflow/python/kernel_tests/template_test.py
+++ b/tensorflow/python/kernel_tests/template_test.py
@@ -21,6 +21,7 @@ import traceback
 
 from tensorflow.python.client import session
 from tensorflow.python.framework import random_seed
+from tensorflow.python.ops import array_ops
 from tensorflow.python.ops import init_ops
 from tensorflow.python.ops import math_ops
 from tensorflow.python.ops import template
@@ -318,6 +319,54 @@ class TemplateTest(test.TestCase):
     tmpl2()
     self.assertEqual(custom_getter_count[0], 2)
 
+  def test_name_scopes_for_variable_scopes(self):
+    # Test that name scopes are not unnecessarily uniquified (but are
+    # still uniquified when necessary).
+    def linear_module(x, output_size):
+      w = variable_scope.get_variable(
+          "w", shape=[x.get_shape()[1], output_size],
+          initializer=init_ops.zeros_initializer())
+      b = variable_scope.get_variable(
+          "b", shape=[output_size],
+          initializer=init_ops.zeros_initializer())
+      return (math_ops.matmul(x, w) + b), w
+
+    def make_linear_module(output_size, name):
+      return template.make_template(
+          name,
+          linear_module,
+          output_size=output_size,
+          create_scope_now_=True)
+
+    inputs = array_ops.ones((3, 4))
+
+    linear1 = make_linear_module(output_size=2, name="foo")
+    outputs_a, w1 = linear1(inputs)
+    outputs_b, _ = linear1(inputs)
+    self.assertEquals("foo", linear1.variable_scope.name)
+    self.assertEquals("foo/w:0", w1.name)
+    self.assertEquals("foo/add:0", outputs_a.name,
+                      "First application of template should get "
+                      "same name scope as variables.")
+    self.assertEquals("foo_1/add:0", outputs_b.name,
+                      "Second application of template should get "
+                      "a freshly uniquified name scope.")
+
+    linear2 = make_linear_module(output_size=2, name="foo")
+    outputs_c, w2 = linear2(inputs)
+    outputs_d, _ = linear2(inputs)
+    self.assertEquals("foo_1", linear2.variable_scope.name,
+                      "New template gets a freshly uniquified variable scope "
+                      "because 'foo' is already taken.")
+    self.assertEquals("foo_1/w:0", w2.name)
+    self.assertEquals("foo_1_1/add:0", outputs_c.name,
+                      "First application of template would get "
+                      "same name scope as variables, but 'foo_1' is already "
+                      "a name scope.")
+    self.assertEquals("foo_1_2/add:0", outputs_d.name,
+                      "Second application of template should also get "
+                      "a freshly uniquified name scope.")
+
 
 if __name__ == "__main__":
   test.main()
diff --git a/tensorflow/python/ops/template.py b/tensorflow/python/ops/template.py
index d4778977b3..80dd74521b 100644
--- a/tensorflow/python/ops/template.py
+++ b/tensorflow/python/ops/template.py
@@ -197,8 +197,9 @@ class Template(object):
     if name is None:
       raise ValueError("name cannot be None.")
     if create_scope_now:
-      with variable_scope.variable_scope(
-          self._unique_name, self._name,
+      with variable_scope._pure_variable_scope(  # pylint:disable=protected-access
+          (self._unique_name or
+           variable_scope._get_unique_variable_scope(self._name)),  # pylint:disable=protected-access
           custom_getter=self._custom_getter) as vs:
         self._variable_scope = vs
     else: