class_docstr: Handles input pipeline for SimdMeshImpl.
class_docstr_tok: ["Handles","input","pipeline","for","SimdMeshImpl","."]
docstr: Input pipeline for the SIMD implementation of MeshTensorflow.
  
  Args:
    simd_mesh_impl: A mtf.simd_mesh_impl.SimdMeshImpl object.
    ds_creator: A function that creates a dataset.
    mtf_input_shapes: A list of mtf.Shape. Then length of it must be equal
      to the number of elements generated by the ds_creator. NOTE, we assume:
        1. The 0-th dimension is the batch dimension.
        2. The batch dimension is consistent across all input shapes in
           mtf_input_shapes.
    ds_prefetch_size: The buffer size for prefetching
      (default tf.data.experimental.AUTOTUNE).
    external_worker: Whether you have an external tpu_worker or not. Set it to
      False if you run the program locally, for example, during local unit
      test.
    is_eval_mode: In evaluation mode, only one dataset object will be created,
      as opposed to one dataset for each sub-batch. Default is False. Set it
      to True during evaluation, to ensure that one evaluation instance will
      be used once and only once.
  Note:
    1. The efficiency is optimized according to the shape of the 0-th tensor:
       mtf_input_shapes[0]. We recommand you to put the largest tensor as the
       0-th input.
    2. You need to call start_infeed_thread() before your train ops.
  Example:
      simd_mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(...)
      # ds_creator is function that creates a tf.data.Dataset.
      # This Dataset must return single examples (no batch dimension).
      def ds_creator():
        return tf.data.Dataset.from_tensors(x)
      # mtf_input_shapes is a list of Shapes of all input tensors given by the
      # dataset. All shapes must begin with the same batch dimension.
      simd_input_reader = SimdMeshImplInputReader(simd_mesh_impl,
                                                  ds_creator,
                                                  mtf_input_shapes)
  
      batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
      io_dim = mtf.Dimension('io', FLAGS.io_size)
      mtf_input_shapes = [mtf.Shape([batch_dim, io_dim])]
  
      infeed_queue = simd_input_reader.infeed_queue
      tpu_train_computation = tpu.replicate(
          computation=model_fn,
          inputs=[[]] * num_cores,
          infeed_queue=infeed_queue, ...)
  
      with tf.Session() as sess:
        simd_input_reader.start_infeed_thread(sess,
                                              number_steps=num_training_steps)
        for _ in range(num_training_steps):
          sess.run(tpu_train_computation)
docstr_tok: ["Input","pipeline","for","the","SIMD","implementation","of","MeshTensorflow",".","Args",":","simd_mesh_impl",":","A","mtf.simd_mesh_impl.SimdMeshImpl","object",".","ds_creator",":","A","function","that","creates","a","dataset",".","mtf_input_shapes",":","A","list","of","mtf.Shape",".","Then","length","of","it","must","be","equal","to","the","number","of","elements","generated","by","the","ds_creator",".","NOTE",",","we","assume",":","1",".","The","0-th","dimension","is","the","batch","dimension",".","2",".","The","batch","dimension","is","consistent","across","all","input","shapes","in","mtf_input_shapes",".","ds_prefetch_size",":","The","buffer","size","for","prefetching","(","default","tf.data.experimental.AUTOTUNE",")",".","external_worker",":","Whether","you","have","an","external","tpu_worker","or","not",".","Set","it","to","False","if","you","run","the","program","locally",",","for","example",",","during","local","unit","test",".","is_eval_mode",":","In","evaluation","mode",",","only","one","dataset","object","will","be","created",",","as","opposed","to","one","dataset","for","each","sub-batch",".","Default","is","False",".","Set","it","to","True","during","evaluation",",","to","ensure","that","one","evaluation","instance","will","be","used","once","and","only","once",".","Note",":","1",".","The","efficiency","is","optimized","according","to","the","shape","of","the","0-th","tensor",":","mtf_input_shapes","[","0","]",".","We","recommand","you","to","put","the","largest","tensor","as","the","0-th","input",".","2",".","You","need","to","call","start_infeed_thread","(",")","before","your","train","ops",".","Example",":","simd_mesh_impl","=","mtf.simd_mesh_impl.SimdMeshImpl","(","...",")","#","ds_creator","is","function","that","creates","a","tf.data.Dataset",".","#","This","Dataset","must","return","single","examples","(","no","batch","dimension",")",".","def","ds_creator","(",")",":","return","tf.data.Dataset.from_tensors","(","x",")","#","mtf_input_shapes","is","a","list","of","Shapes","of","all","input","tensors","given","by","the","#","dataset",".","All","shapes","must","begin","with","the","same","batch","dimension",".","simd_input_reader","=","SimdMeshImplInputReader","(","simd_mesh_impl",",","ds_creator",",","mtf_input_shapes",")","batch_dim","=","mtf.Dimension","(","'batch","'",",","FLAGS.batch_size",")","io_dim","=","mtf.Dimension","(","'io","'",",","FLAGS.io_size",")","mtf_input_shapes","=","[","mtf.Shape","(","[","batch_dim",",","io_dim","]",")","]","infeed_queue","=","simd_input_reader.infeed_queue","tpu_train_computation","=","tpu.replicate","(","computation=model_fn",",","inputs=","[","[","]","]","*","num_cores",",","infeed_queue=infeed_queue",",","...",")","with","tf.Session","(",")","as","sess",":","simd_input_reader.start_infeed_thread","(","sess",",","number_steps=num_training_steps",")","for","_","in","range","(","num_training_steps",")",":","sess.run","(","tpu_train_computation",")"]
code: 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,
  ):
      """Input pipeline for the SIMD implementation of MeshTensorflow.
  
      Args:
        simd_mesh_impl: A mtf.simd_mesh_impl.SimdMeshImpl object.
        ds_creator: A function that creates a dataset.
        mtf_input_shapes: A list of mtf.Shape. Then length of it must be equal
          to the number of elements generated by the ds_creator. NOTE, we assume:
            1. The 0-th dimension is the batch dimension.
            2. The batch dimension is consistent across all input shapes in
               mtf_input_shapes.
        ds_prefetch_size: The buffer size for prefetching
          (default tf.data.experimental.AUTOTUNE).
        external_worker: Whether you have an external tpu_worker or not. Set it to
          False if you run the program locally, for example, during local unit
          test.
        is_eval_mode: In evaluation mode, only one dataset object will be created,
          as opposed to one dataset for each sub-batch. Default is False. Set it
          to True during evaluation, to ensure that one evaluation instance will
          be used once and only once.
      Note:
        1. The efficiency is optimized according to the shape of the 0-th tensor:
           mtf_input_shapes[0]. We recommand you to put the largest tensor as the
           0-th input.
        2. You need to call start_infeed_thread() before your train ops.
      Example:
          simd_mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(...)
          # ds_creator is function that creates a tf.data.Dataset.
          # This Dataset must return single examples (no batch dimension).
          def ds_creator():
            return tf.data.Dataset.from_tensors(x)
          # mtf_input_shapes is a list of Shapes of all input tensors given by the
          # dataset. All shapes must begin with the same batch dimension.
          simd_input_reader = SimdMeshImplInputReader(simd_mesh_impl,
                                                      ds_creator,
                                                      mtf_input_shapes)
  
          batch_dim = mtf.Dimension('batch', FLAGS.batch_size)
          io_dim = mtf.Dimension('io', FLAGS.io_size)
          mtf_input_shapes = [mtf.Shape([batch_dim, io_dim])]
  
          infeed_queue = simd_input_reader.infeed_queue
          tpu_train_computation = tpu.replicate(
              computation=model_fn,
              inputs=[[]] * num_cores,
              infeed_queue=infeed_queue, ...)
  
          with tf.Session() as sess:
            simd_input_reader.start_infeed_thread(sess,
                                                  number_steps=num_training_steps)
            for _ in range(num_training_steps):
              sess.run(tpu_train_computation)
      """
      super(SimdMeshImplInputReader, self).__init__()
      assert mtf_input_shapes
      assert isinstance(mtf_input_shapes, list)
  
      # TODO(lehou): Support nested structures for ds_creator, mtf_input_shapes.
      self._simd_mesh_impl = simd_mesh_impl
      self._p_dev = ProcessDevices(simd_mesh_impl)
      self._ds_creator = ds_creator
      self._mtf_input_shapes = mtf_input_shapes
      self._ds_prefetch_size = ds_prefetch_size
      self._external_worker = external_worker
      self._is_eval_mode = is_eval_mode
  
      self._gen_infeed_queue()
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import collections
  import threading
  import time
  import numpy
  from six.moves import range
  from six.moves import zip
  import tensorflow.compat.v1
  from tensorflow.python.framework import ops
  from tensorflow.python.tpu import tpu_feed
label_desc: Use image node_id 1 to create a new SimdMeshImplInputReader object from inherited base classes: object with example: obj = SimdMeshImplInputReader(simd_mesh_impl, ds_creator, mtf_input_shapes, ds_prefetch_size, external_worker, is_eval_mode)