(ctx, model_fn_wrapper, dequeue_fn)
return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get()
Executes `model_fn_wrapper` multiple times on all TPU shards.
Executes `model_fn_wrapper` multiple times on all TPU shards.
Executes model_fn_wrapper multiple times on all TPU shards .
def _eval_on_tpu_system(ctx, model_fn_wrapper, dequeue_fn):
  """Executes `model_fn_wrapper` multiple times on all TPU shards."""
  iterations_per_loop_var = _create_or_get_iterations_per_loop()

  (single_tpu_eval_step, host_calls, captured_scaffold_fn, captured_eval_hooks
  ) = model_fn_wrapper.convert_to_single_tpu_eval_step(dequeue_fn)

  def multi_tpu_eval_steps_on_single_shard():
    loop_vars = [_ZERO_LOSS]
    if model_fn_wrapper._eval_cache_fn is not None:
      batch_size = ctx.global_batch_size
      num_shards = ctx._config._tpu_config.num_shards
      loop_vars += model_fn_wrapper._eval_cache_fn(batch_size // num_shards)

    return training_loop.repeat(
        iterations_per_loop_var,
        single_tpu_eval_step,
        loop_vars)

  compile_op, ret = tpu.split_compile_and_shard(
      multi_tpu_eval_steps_on_single_shard,
      inputs=[],
      num_shards=ctx.num_replicas,
      outputs_from_all_shards=False,
      device_assignment=ctx.device_assignment)

  loss = ret[0]
  scaffold = _get_scaffold(captured_scaffold_fn)
  return compile_op, loss, host_calls, scaffold, captured_eval_hooks.get()
def _eval_on_tpu_system ( ctx , model_fn_wrapper , dequeue_fn ) : iterations_per_loop_var = _create_or_get_iterations_per_loop ( ) ( single_tpu_eval_step , host_calls , captured_scaffold_fn , captured_eval_hooks ) = model_fn_wrapper . convert_to_single_tpu_eval_step ( dequeue_fn ) def multi_tpu_eval_steps_on_single_shard ( ) : loop_vars = [ _ZERO_LOSS ] if model_fn_wrapper . _eval_cache_fn is not None : batch_size = ctx . global_batch_size num_shards = ctx . _config . _tpu_config . num_shards loop_vars += model_fn_wrapper . _eval_cache_fn ( batch_size // num_shards ) return training_loop . repeat ( iterations_per_loop_var , single_tpu_eval_step , loop_vars ) compile_op , ret = tpu . split_compile_and_shard ( multi_tpu_eval_steps_on_single_shard , inputs = [ ] , num_shards = ctx . num_replicas , outputs_from_all_shards = False , device_assignment = ctx . device_assignment ) loss = ret [ 0 ] scaffold = _get_scaffold ( captured_scaffold_fn ) return compile_op , loss , host_calls , scaffold , captured_eval_hooks . get ( )