def computation():
      """Compute tpu tensors used in export_outputs.

      Passed to rewrite_for_inference so that model_fn will be called under
      the rewriting contexts. Only tpu tensors are returned, but export_outputs
      and scaffold are captured.

      Returns:
         A list of Tensors used in export_outputs and not marked for
         outside_compilation.
      """
      # We should only call model fn once and it should be inside `computation`
      # so that building the graph will happen under `rewrite_for_inference`.
      mode = model_fn_lib.ModeKeys.PREDICT
      estimator_spec = self._call_model_fn(features, labels, mode, config)

      # We pick the TPU tensors out from `export_output` and later return them
      # from `computation` for rewriting.
      tensors_dict = collections.OrderedDict(
          (k, _export_output_to_tensors(v))
          for k, v in six.iteritems(estimator_spec.export_outputs))
      tensors = nest.flatten(tensors_dict)
      tpu_tensors = [t for t in tensors if t is not None]

      # We cannot return anything other than `tpu_tensors` here so we capture
      # the rest for later use.
      capture.capture((estimator_spec, tensors_dict, tensors))
      return tpu_tensors
Compute tpu tensors used in export_outputs.

      Passed to rewrite_for_inference so that model_fn will be called under
      the rewriting contexts. Only tpu tensors are returned, but export_outputs
      and scaffold are captured.

      Returns:
         A list of Tensors used in export_outputs and not marked for
         outside_compilation.
      
src/tpu_estimator.py