Make TpuContext thread-local.

tf.contrib.AsyncCheckpointSaverHook creates a thread to asynchronously run the export, which currently breaks when exporting a TPU graph.

PiperOrigin-RevId: 258184145

tensorflow/python/tpu/tpu_function.py

@@ -20,9 +20,10 @@ from __future__ import division
 from __future__ import print_function
 
 import contextlib
+import threading
 
 
-class TpuContext(object):
+class TpuContext(threading.local):
   """A context object holding state about the TPU computation being built."""
 
   def __init__(self):