Fixed the issue that each invocation of model.fit/evaluate/predict modifies the (#23280)

tatatodd · web-flow · commit 748435b8ef55 · 2018-10-25T20:26:00.000-07:00
graph.

PiperOrigin-RevId: 218793646
diff --git a/tensorflow/contrib/tpu/python/tpu/keras_support.py b/tensorflow/contrib/tpu/python/tpu/keras_support.py
@@ -97,14 +97,25 @@
 
 # TODO(b/114775106): temporary shim to optionally initialize the TPU
 # This increases the odds our session is initialized, but shouldn't be needed.
+_TEST_REWRITE_OP = None
+
+
 def _maybe_initialize_tpu(session):
   """Initialize the TPU if it has not already been initialized."""
+  global _TEST_REWRITE_OP
   try:
+    # Try to use cached version to avoid another ground of graph optimization.
+    test_rewrite_op = _TEST_REWRITE_OP
+    if (test_rewrite_op is None or
+        test_rewrite_op[0].graph != ops.get_default_graph()):
+
+      def test_op():
+        return constant_op.constant(1) + constant_op.constant(1)
 
-    def test_op():
-      return constant_op.constant(1) + constant_op.constant(1)
+      test_rewrite_op = tpu.rewrite(test_op)
+      _TEST_REWRITE_OP = test_rewrite_op
 
-    session.run(tpu.rewrite(test_op))
+    session.run(test_rewrite_op)
   except errors.FailedPreconditionError as _:
     session.run(tpu.initialize_system())
 
