class_docstr: Handles input pipeline for SimdMeshImpl.
class_docstr_tok: ["Handles","input","pipeline","for","SimdMeshImpl","."]
docstr: Generates _infeed_queue, _enqueue_ops, _input_initializers.
docstr_tok: ["Generates","_infeed_queue",",","_enqueue_ops",",","_input_initializers","."]
code: def _gen_infeed_queue(self):
      """Generates _infeed_queue, _enqueue_ops, _input_initializers."""
      pnum_maps = []
      batch_size = self._mtf_input_shapes[0].to_integer_list[0]
      for mtf_shape in self._mtf_input_shapes:
          # Make sure that the batch size is the same across all input tensors.
          assert batch_size == mtf_shape.to_integer_list[0]
          pnum_maps.append(self._get_pnum_map(mtf_shape))
  
      # For each sub-batch, we need to know which host should read it.
      if self._is_eval_mode:
          # There should be just one dataset-holding host. Make the last host do it.
          hosts_to_hold_ds = [self._p_dev.num_hosts - 1]
      else:
          hosts_to_hold_ds = self._get_hosts_to_hold_ds(pnum_maps[0])
      sub_batch_size = batch_size // len(hosts_to_hold_ds)
      tf.logging.info("MTF sub_batch_size: {}".format(sub_batch_size))
      assert sub_batch_size * len(hosts_to_hold_ds) == batch_size
  
      def sub_batch_ds_creator():
          return (
              self._ds_creator()
              .batch(sub_batch_size, drop_remainder=True)
              .prefetch(self._ds_prefetch_size)
          )
  
      sub_batch_slicer_list = []
      # For each sub-batch, create a SubBatchSlicer object.
      for sub_batch_i, host_id in enumerate(hosts_to_hold_ds):
          # Get the list of pnums for each input.
          if self._is_eval_mode:
              all_sub_batch_pnums = [
                  pnum_map.flatten().tolist() for pnum_map in pnum_maps
              ]
  
              sub_batch_slicer_list.append(
                  SubBatchSlicer(
                      sub_batch_ds_creator,
                      host_id,
                      all_sub_batch_pnums,
                      self._simd_mesh_impl,
                      self._mtf_input_shapes,
                      self._external_worker,
                      global_batch=True,
                  )
              )
          else:
              all_sub_batch_pnums = []
              for pnum_map in pnum_maps:
                  sub_batch_pnums = (
                      pnum_map[sub_batch_i, ...].flatten().tolist()
                  )
                  all_sub_batch_pnums.append(sub_batch_pnums)
  
              sub_batch_slicer_list.append(
                  SubBatchSlicer(
                      sub_batch_ds_creator,
                      host_id,
                      all_sub_batch_pnums,
                      self._simd_mesh_impl,
                      self._mtf_input_shapes,
                      self._external_worker,
                      global_batch=False,
                  )
              )
  
      # Slots for all laidout tensors.
      all_laidout_tensors = [
          [_NO_DATA] * len(self._mtf_input_shapes)
          for _ in range(self._p_dev.num_cores)
      ]
  
      # Read tf_tensors, put them in slots.
      for sub_batch_slicer in sub_batch_slicer_list:
          for tf_tensor, pnum, input_i in sub_batch_slicer.get_slices():
              all_laidout_tensors[pnum][input_i] = tf_tensor
  
      # Make sure that there are no Nones in all_laidout_tensors.
      for laidout_tensors in all_laidout_tensors:
          assert _NO_DATA not in laidout_tensors
  
      with ops.device(
          _host_id_to_tf_device(
              hosts_to_hold_ds[0], self._external_worker
          )
      ):
          (
              self._infeed_queue,
              self._enqueue_ops,
          ) = self._enqueue_laidout_tensors(all_laidout_tensors)
  
      self._input_initializers = [
          s.initializer for s in sub_batch_slicer_list
      ]
  
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 4 for calling the SimdMeshImplInputReader obj's underlying member method code with example usage: obj._gen_infeed_queue() and returns: self