def _train_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_train_step, host_call, captured_scaffold_fn,
   captured_training_hooks) = (
       model_fn_wrapper.convert_to_single_tpu_train_step(dequeue_fn))

  def multi_tpu_train_steps_on_single_shard():
    loop_vars = [_INITIAL_LOSS]
    if model_fn_wrapper._train_cache_fn is not None:
      batch_size = ctx.global_batch_size
      num_shards = ctx._config._tpu_config.num_shards
      loop_vars += model_fn_wrapper._train_cache_fn(batch_size // num_shards)

    return training_loop.repeat(
        iterations_per_loop_var,
        single_tpu_train_step,
        loop_vars)

  compile_op, ret = tpu.split_compile_and_shard(
      multi_tpu_train_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_call, scaffold, captured_training_hooks.get()
Executes `model_fn_wrapper` multiple times on all TPU shards.
src/tpu_estimator.py