Added fail_if_exists as an argument for write_to_variables.

PiperOrigin-RevId: 277081924

Author: mn-robot

Diff for: morph_net/framework/tpu_util.py

@@ -94,20 +94,21 @@ def maybe_convert_to_variable(tensor):
 var_store = {}
 
 
-def write_to_variable(tensor):
+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]
 
-  # 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)
+  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):
     variable = tf.get_variable(