aboutsummaryrefslogtreecommitdiffhomepage
path: root/tensorflow/contrib/tpu/__init__.py
diff options
context:
space:
mode:
authorGravatar Russell Power <power@google.com>2018-09-28 16:41:58 -0700
committerGravatar TensorFlower Gardener <gardener@tensorflow.org>2018-09-28 16:45:25 -0700
commit541677bfee008a093daab2d033bd72650d886126 (patch)
tree0b8373dfac0b18ac70f29f78e379ceeb8b185015 /tensorflow/contrib/tpu/__init__.py
parent0a341bbcb35d72d14bfda17f7f0cb0c61f323bce (diff)
Add option to disable initialization/shutdown of the TPU.
PiperOrigin-RevId: 215016286
Diffstat (limited to 'tensorflow/contrib/tpu/__init__.py')
-rw-r--r--tensorflow/contrib/tpu/__init__.py3
1 files changed, 3 insertions, 0 deletions
diff --git a/tensorflow/contrib/tpu/__init__.py b/tensorflow/contrib/tpu/__init__.py
index 766466968a..6ce6b779a2 100644
--- a/tensorflow/contrib/tpu/__init__.py
+++ b/tensorflow/contrib/tpu/__init__.py
@@ -55,7 +55,9 @@
@@TPUDistributionStrategy
@@keras_to_tpu_model
+
@@AsyncCheckpointSaverHook
+@@TPUInMemoryEvalHook
"""
from __future__ import absolute_import
@@ -65,6 +67,7 @@ from __future__ import print_function
# pylint: disable=wildcard-import,unused-import
from tensorflow.contrib.tpu.python import profiler
from tensorflow.contrib.tpu.python.ops.tpu_ops import *
+from tensorflow.contrib.tpu.python.tpu.async_checkpoint import *
from tensorflow.contrib.tpu.python.tpu.bfloat16 import *
from tensorflow.contrib.tpu.python.tpu.device_assignment import *
from tensorflow.contrib.tpu.python.tpu.keras_support import tpu_model as keras_to_tpu_model