Make TPU summaries work with non-variable BatchNorm gammas (and other non-variable tensors).

PiperOrigin-RevId: 262224145

Author: shraman-rc

morph_net/framework/tpu_util.py

@@ -58,6 +58,8 @@ def maybe_convert_to_variable(tensor):
     return the original input tensor.
   """
   op = tensor.op
+  if is_on_cpu() and tensor in var_store:
+    return var_store[tensor]
   if op.type != 'ReadVariableOp':
     # No need to convert.
     return tensor
@@ -102,7 +104,9 @@ def write_to_variable(tensor):
         use_resource=True)
   var_store[tensor] = variable
   with tf.control_dependencies([variable.assign(tensor)]):
-    return tf.identity(tensor)
+    tensor_copy = tf.identity(tensor)
+  var_store[tensor_copy] = variable
+  return tensor_copy
 
 
 def read_from_variable(tensor):
@@ -113,3 +117,8 @@ def read_from_variable(tensor):
   else:
     # Current read, but only works on TPU.
     return tensor
+
+
+def is_intermediate_var(v):
+  """Returns True if `v` was created by `write_to_variable` above."""
+  return v in var_store.values()