Add custom_getter to variable scope.

Change: 128486774

Author: ebrevdo

tensorflow/python/kernel_tests/variable_scope_test.py

@@ -673,5 +673,54 @@ def _part_axis_1(**unused_kwargs):
     self.assertEqual(n1_0.get_shape(), (2, 1, 2))
     self.assertEqual(n1_1.get_shape(), (2, 1, 2))
 
+
+class VariableScopeWithCustomGetterTest(tf.test.TestCase):
+
+  def testNonCallableGetterFails(self):
+    with self.assertRaisesRegexp(ValueError, r"custom_getter .* not callable:"):
+      with tf.variable_scope("scope0", custom_getter=3):
+        tf.get_variable("name0")
+    with self.assertRaisesRegexp(ValueError, r"custom_getter .* not callable:"):
+      tf.get_variable("name0", custom_getter=3)
+
+  def testNoSideEffectsWithIdentityCustomGetter(self):
+    called = [0]
+    def custom_getter(getter, *args, **kwargs):
+      called[0] += 1
+      return getter(*args, **kwargs)
+    with tf.variable_scope("scope", custom_getter=custom_getter) as scope:
+      v = tf.get_variable("v", [1])
+    with tf.variable_scope(scope, reuse=True):
+      v2 = tf.get_variable("v", [1])
+    with tf.variable_scope("new_scope") as new_scope:
+      v3 = tf.get_variable("v3", [1])
+    with tf.variable_scope(new_scope, reuse=True, custom_getter=custom_getter):
+      v4 = tf.get_variable("v3", [1])
+
+    self.assertEqual(v, v2)
+    self.assertEqual(v3, v4)
+    self.assertEqual(3, called[0])  # skipped one in the first new_scope
+
+  def testGetterThatCreatesTwoVariablesAndSumsThem(self):
+    def custom_getter(getter, name, *args, **kwargs):
+      g_0 = getter("%s/0" % name, *args, **kwargs)
+      g_1 = getter("%s/1" % name, *args, **kwargs)
+      with tf.name_scope("custom_getter"):
+        return g_0 + g_1
+
+    with tf.variable_scope("scope", custom_getter=custom_getter):
+      v = tf.get_variable("v", [1, 2, 3])
+
+    self.assertEqual([1, 2, 3], v.get_shape())
+    true_vars = tf.trainable_variables()
+    self.assertEqual(2, len(true_vars))
+    self.assertEqual("scope/v/0:0", true_vars[0].name)
+    self.assertEqual("scope/v/1:0", true_vars[1].name)
+    self.assertEqual("custom_getter/add:0", v.name)
+    with self.test_session() as sess:
+      tf.initialize_all_variables().run()
+      np_vars, np_v = sess.run([true_vars, v])
+      self.assertAllClose(np_v, sum(np_vars))
+
 if __name__ == "__main__":
   tf.test.main()