message: Use outside compilation in TPUEstimator.export_savedmodel().
diff: @@ -17,7 +17,6 @@ from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  -import contextlib
  import math
  from tensorflow.contrib.tpu.python.tpu import tpu
  @@ -279,11 +278,10 @@ class _TPUEmbeddingColumn(_TPUBaseEmbeddingColumn, fc._EmbeddingColumn):
  def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None):
  if tpu.under_tpu_inference_context():
  - # TODO(shizhiw, b/112012627, b/112336539): Replace _outside_all_rewrites()
  - # with outside compilation.
  - with _outside_all_rewrites():
  + def host_computation():
  return fc._EmbeddingColumn._get_dense_tensor(
  self, inputs, weight_collections, trainable)
  + return tpu.outside_compilation(host_computation)
  if _is_running_on_cpu():
  return fc._EmbeddingColumn._get_dense_tensor(
  @@ -300,13 +298,6 @@ class _TPUEmbeddingColumn(_TPUBaseEmbeddingColumn, fc._EmbeddingColumn):
  return tensor
  -@contextlib.contextmanager
  -def _outside_all_rewrites():
  - """'Break out' of a tpu.rewrite() (or shard(), etc.)."""
  - with ops.control_dependencies(None):
  - yield
  -
  -
  class _TPUSharedEmbeddingColumn(_TPUBaseEmbeddingColumn,
  fc._SharedEmbeddingColumn):
  """Core Shared Embedding Column."""
  @@ -385,11 +376,10 @@ class _TPUSharedEmbeddingColumn(_TPUBaseEmbeddingColumn,
  def _get_dense_tensor(self, inputs, weight_collections=None, trainable=None):
  if tpu.under_tpu_inference_context():
  - # TODO(shizhiw, b/112012627, b/112336539): Replace _outside_all_rewrites()
  - # with outside compilation.
  - with _outside_all_rewrites():
  + def host_computation():
  return fc._SharedEmbeddingColumn._get_dense_tensor(
  self, inputs, weight_collections, trainable)
  + return tpu.outside_compilation(host_computation)
  if _is_running_on_cpu():
  return fc._SharedEmbeddingColumn._get_dense_tensor(
  