message: Use tpu.rewrite instead of tpu.rewrite_for_inference with PartitionedCallOp in TPUEstimator.py.
diff: @@ -1363,7 +1363,7 @@ def call_computation(computation,
  # TPU core with every `Session.run()` call. Note that the entire inference
  # graph executes on a single core, and that invocations of this graph
  # will round-robin among the cores attached to a host.
  - @function.Defun()
  + @function.Defun(capture_resource_var_by_value=False)
  def tpu_subgraph():
  return computation()
  @@ -2464,8 +2464,14 @@ class TPUEstimator(estimator_lib.Estimator):
  device_assignment = ctx.device_assignment
  else:
  device_assignment = None
  +
  + if self._experimental_exported_model_uses_all_cores:
  + tensors_on_cpu = tpu.rewrite(
  + tpu_computation, device_assignment=device_assignment)
  + else:
  tensors_on_cpu = tpu.rewrite_for_inference(
  tpu_computation, device_assignment=device_assignment)
  +
  (estimator_spec, export_outputs_dict, export_outputs_list,
  predictions_dict) = (
  tpu_capture.get())
  