# Path: mesh_tensorflow/experimental/input_reader.py
# _NONE_PNUM = None
# _NO_DATA = None
# def _host_device_to_id(device_str):
# def _host_id_to_tf_device(host_id, external_worker):
#   def __init__(self, sub_batch_ds_creator, host_id, all_sub_batch_pnums,
#                simd_mesh_impl, mtf_input_shapes, external_worker, global_batch):
#   def initializer(self):
#   def get_slices(self):
#   def _validate_args(self):
#   def _init_slice_cache(self):
#   def _slice_tensor(self, input_tensor, mtf_input_shape, pnum):
#   def __init__(self, simd_mesh_impl):
#   def ordered_ordinals(self):
#   def ordered_tpus(self):
#   def ordered_hosts(self):
#   def ordered_host_ids(self):
#   def host_id_to_its_pnums(self):
#   def num_cores(self):
#   def num_hosts(self):
#   def num_cores_per_host(self):
#   def __init__(self,
#                simd_mesh_impl,
#                ds_creator,
#                mtf_input_shapes,
#                ds_prefetch_size=tf.data.experimental.AUTOTUNE,
#                external_worker=True,
#                is_eval_mode=False):
#   def infeed_queue(self):
#   def start_infeed_thread(self, sess, number_steps=-1, initial_wait_sec=0.5):
#     def _thread_fn():
#   def _gen_infeed_queue(self):
#     def sub_batch_ds_creator():
#   def _get_pnum_map(self, mtf_shape):
#   def _get_hosts_to_hold_ds(self, pnum_map):
#     def _get_num_pnums_per_host(sub_batch_pnum_map):
#     def _find_host_id_with_most_pnums_and_least_ds(num_pnums_per_host,
#                                                    num_dss_per_host):
#   def _enqueue_laidout_tensors(self, all_laidout_tensors):
#     def _tpu_ordinal_function_impl(pnum):
#     def _placement_function_impl(pnum):
#   def __init__(self,
#                placement_mesh_impl,
#                ds_creator,
#                mtf_input_shapes,
#                ds_prefetch_size=tf.data.experimental.AUTOTUNE,
#                is_eval_mode=False):
#   def initialize(self, sess):
#   def gpu_placement(self, model_fn):
# class SubBatchSlicer(object):
# class ProcessDevices(object):
# class SimdMeshImplInputReader(object):
# class PlacementMeshImplInputReader(object):
# 
# Path: mesh_tensorflow/experimental/unet.py
# FLAGS = flags.FLAGS
# def get_layout():
# def get_dataset_creator(dataset_str):
#   def _dataset_creator():
#     def _get_stacked_2d_slices(image_3d, label_3d):
#     def _parser_fn(serialized_example):
# def get_input_mtf_shapes(dataset_str):
#   def __init__(self):
#   def record(self, results, pred_output_dir):
#   def finish(self):
#   def _reshape_to_cubes(self, data):
# def conv_with_spatial_partition(
#     x, sampled_2d_slices, image_nx_dim, image_ny_dim, n_filters,
#     keep_p, with_batch_norm, is_training, odim_name, variable_dtype, name):
# def deconv_with_spatial_partition(
#     x, sampled_2d_slices, image_nx_dim, image_ny_dim, n_filters, keep_p,
#     is_training,
#     odim_name, variable_dtype, name):
# def unet_with_spatial_partition(mesh, mesh_impl, dataset_str, images, labels):
#   def scalar(v, dtype):
# class PostProcessor(object):


import re
import mesh_tensorflow as mtf
import numpy as np
import six
import tensorflow.compat.v1 as tf
from six.moves import range
from six.moves import zip
from mesh_tensorflow.experimental import input_reader
from mesh_tensorflow.experimental import unet
from tensorflow.contrib import summary as contrib_summary
from tensorflow.contrib import tpu
from tensorflow.contrib.tpu.python.tpu import device_assignment
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.platform import flags
from tensorflow.python.tpu.ops import tpu_ops
        simd_input_reader.start_infeed_thread(
            sess, FLAGS.num_train_iterations_per_loop)
      else:
        placement_input_reader.initialize(sess)

      for step in range(FLAGS.num_train_iterations_per_loop):
        sess.run(train_computation)
        if FLAGS.write_summary:
          sess.run(flush_summary)
        tf.logging.info('train steps: {}'.format(step))

  with tf.Graph().as_default():
    if FLAGS.write_summary:
      summary_writer = contrib_summary.create_file_writer(FLAGS.summary_dir)
      with summary_writer.as_default(), (
          contrib_summary.always_record_summaries()):
        _run_train_phase()
    else:
      _run_train_phase()


def _eval_phase(mesh_context, config, master):
  """Handles input pipeline and evaluates the network."""
  if FLAGS.num_eval_iterations_per_loop <= 0:
    return

  def _run_eval_phase():
    """The real function that runs the evaluation phase."""
    # Setup input pipeline.
    ds_creator = unet.get_dataset_creator('eval')

# Path: mesh_tensorflow/experimental/input_reader.py
# _NONE_PNUM = None
# _NO_DATA = None
# def _host_device_to_id(device_str):
# def _host_id_to_tf_device(host_id, external_worker):
#   def __init__(self, sub_batch_ds_creator, host_id, all_sub_batch_pnums,
#                simd_mesh_impl, mtf_input_shapes, external_worker, global_batch):
#   def initializer(self):
#   def get_slices(self):
#   def _validate_args(self):
#   def _init_slice_cache(self):
#   def _slice_tensor(self, input_tensor, mtf_input_shape, pnum):
#   def __init__(self, simd_mesh_impl):
#   def ordered_ordinals(self):
#   def ordered_tpus(self):
#   def ordered_hosts(self):
#   def ordered_host_ids(self):
#   def host_id_to_its_pnums(self):
#   def num_cores(self):
#   def num_hosts(self):
#   def num_cores_per_host(self):
#   def __init__(self,
#                simd_mesh_impl,
#                ds_creator,
#                mtf_input_shapes,
#                ds_prefetch_size=tf.data.experimental.AUTOTUNE,
#                external_worker=True,
#                is_eval_mode=False):
#   def infeed_queue(self):
#   def start_infeed_thread(self, sess, number_steps=-1, initial_wait_sec=0.5):
#     def _thread_fn():
#   def _gen_infeed_queue(self):
#     def sub_batch_ds_creator():
#   def _get_pnum_map(self, mtf_shape):
#   def _get_hosts_to_hold_ds(self, pnum_map):
#     def _get_num_pnums_per_host(sub_batch_pnum_map):
#     def _find_host_id_with_most_pnums_and_least_ds(num_pnums_per_host,
#                                                    num_dss_per_host):
#   def _enqueue_laidout_tensors(self, all_laidout_tensors):
#     def _tpu_ordinal_function_impl(pnum):
#     def _placement_function_impl(pnum):
#   def __init__(self,
#                placement_mesh_impl,
#                ds_creator,
#                mtf_input_shapes,
#                ds_prefetch_size=tf.data.experimental.AUTOTUNE,
#                is_eval_mode=False):
#   def initialize(self, sess):
#   def gpu_placement(self, model_fn):
# class SubBatchSlicer(object):
# class ProcessDevices(object):
# class SimdMeshImplInputReader(object):
# class PlacementMeshImplInputReader(object):
# 
# Path: mesh_tensorflow/experimental/unet.py
# FLAGS = flags.FLAGS
# def get_layout():
# def get_dataset_creator(dataset_str):
#   def _dataset_creator():
#     def _get_stacked_2d_slices(image_3d, label_3d):
#     def _parser_fn(serialized_example):
# def get_input_mtf_shapes(dataset_str):
#   def __init__(self):
#   def record(self, results, pred_output_dir):
#   def finish(self):
#   def _reshape_to_cubes(self, data):
# def conv_with_spatial_partition(
#     x, sampled_2d_slices, image_nx_dim, image_ny_dim, n_filters,
#     keep_p, with_batch_norm, is_training, odim_name, variable_dtype, name):
# def deconv_with_spatial_partition(
#     x, sampled_2d_slices, image_nx_dim, image_ny_dim, n_filters, keep_p,
#     is_training,
#     odim_name, variable_dtype, name):
# def unet_with_spatial_partition(mesh, mesh_impl, dataset_str, images, labels):
#   def scalar(v, dtype):
# class PostProcessor(object):


# Path: mesh_tensorflow/experimental/model_executor.py
import re
import mesh_tensorflow as mtf
import numpy as np
import six
import tensorflow.compat.v1 as tf
from six.moves import range
from six.moves import zip
from mesh_tensorflow.experimental import input_reader
from mesh_tensorflow.experimental import unet
from tensorflow.contrib import summary as contrib_summary
from tensorflow.contrib import tpu
from tensorflow.contrib.tpu.python.tpu import device_assignment
from tensorflow.python.framework import ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.platform import flags
from tensorflow.python.tpu.ops import tpu_ops

        simd_input_reader.start_infeed_thread(
            sess, FLAGS.num_train_iterations_per_loop)
      else:
        placement_input_reader.initialize(sess)

      for step in range(FLAGS.num_train_iterations_per_loop):
        sess.run(train_computation)
        if FLAGS.write_summary:
          sess.run(flush_summary)
        tf.logging.info('train steps: {}'.format(step))

  with tf.Graph().as_default():
    if FLAGS.write_summary:
      summary_writer = contrib_summary.create_file_writer(FLAGS.summary_dir)
      with summary_writer.as_default(), (
          contrib_summary.always_record_summaries()):
        _run_train_phase()
    else:
      _run_train_phase()


def _eval_phase(mesh_context, config, master):
  """Handles input pipeline and evaluates the network."""
  if FLAGS.num_eval_iterations_per_loop <= 0:
    return

  def _run_eval_phase():
    """The real function that runs the evaluation phase."""
    # Setup input pipeline.
    ds_creator = unet.get_dataset_creator('eval')

    mtf_shapes = unet.get_input_mtf_shapes('eval')