Make tpu_util.write_to_variable compatible with both Keras models and models that use variable_scope.

PiperOrigin-RevId: 302173658

Author: shraman-rc

morph_net/framework/tpu_util.py

@@ -92,31 +92,23 @@ def maybe_convert_to_variable(tensor):
 
 
 var_store = {}
+top_level_scope = tf.get_variable_scope()
 
 
 def write_to_variable(tensor, fail_if_exists=True):
   """Saves a tensor for later retrieval on CPU."""
   # Only relevant for debugging.
-  debug_name = 'tpu_util__' + tensor.name.split(':')[0].split('/')[-1]
-
-  if fail_if_exists:
-    # Note: reuse cannot be changed from True to False, so we just check if
-    # the variable exists.
-    with tf.variable_scope('', reuse=True):
-      try:
-        tf.get_variable(debug_name)
-      except ValueError:
-        pass  # Variable with name=debug_name does not exist; proceed.
-      else:
-        raise ValueError('Variable %s already exists!' % debug_name)
-
-  with tf.variable_scope('', reuse=tf.compat.v1.AUTO_REUSE):
+  debug_name = 'tpu_util__' + tensor.name.split(':')[0]
+
+  reuse = False if fail_if_exists else tf.compat.v1.AUTO_REUSE
+  with tf.variable_scope(top_level_scope, reuse=reuse):
     variable = tf.get_variable(
         name=debug_name,
         shape=tensor.shape,
         dtype=tensor.dtype,
         trainable=False,
         use_resource=True)
+
   var_store[tensor] = variable
   with tf.control_dependencies([variable.assign(tensor)]):
     tensor_copy = tf.identity(tensor)

morph_net/framework/tpu_util_test.py

@@ -79,15 +79,30 @@ def test_noop(self):
     self.assertEqual(tpu_util.maybe_convert_to_variable(relu), relu)
 
   def test_write_to_variable(self):
-    with tf.variable_scope(''):
-      foo = tf.constant(0.)
-      tpu_util.write_to_variable(foo)
+    foo = tf.constant(0., name='foo')
+    tpu_util.write_to_variable(foo)
+
+    with self.assertRaises(ValueError):
+      tpu_util.write_to_variable(foo, fail_if_exists=True)
+
+    # Variable sharing behavior should be dictated by `fail_if_exists` which
+    # overrides the effect of outer scopes.
+    with tf.variable_scope('', reuse=True):
+      # should fail to return existing variable even though reuse=True
+      with self.assertRaises(ValueError):
+        tpu_util.write_to_variable(foo, fail_if_exists=True)
+
+    with tf.variable_scope('', reuse=False):
+      # should return existing variable even though reuse=False
+      foo_copy = tpu_util.write_to_variable(foo, fail_if_exists=False)
+      self.assertEqual(tpu_util.var_store[foo], tpu_util.var_store[foo_copy])
+      self.assertLen(set(tpu_util.var_store.values()), 1)
+
     with tf.variable_scope('', reuse=True):
-      bar = tf.constant(0.)
+      # should create new variable even though reuse=True
+      bar = tf.constant(0., name='bar')
       tpu_util.write_to_variable(bar)
-    with tf.variable_scope('', reuse=tf.compat.v1.AUTO_REUSE):
-      zee = tf.constant(0.)
-      tpu_util.write_to_variable(zee)
+      self.assertLen(set(tpu_util.var_store.values()), 2)
 
 if __name__ == '__main__':
   tf.test.main()