class_docstr: Reads and distributes a sub-batch on a host.
class_docstr_tok: ["Reads","and","distributes","a","sub-batch","on","a","host","."]
docstr: Yields sliced tensors and which remote pnums they should go to.
  
  Yields:
    tf_tensor: The sliced tensor.
    pnum: Which process number the tf_tensor should to go.
    input_i: The input ordinal of the tf_tensor.
docstr_tok: ["Yields","sliced","tensors","and","which","remote","pnums","they","should","go","to",".","Yields",":","tf_tensor",":","The","sliced","tensor",".","pnum",":","Which","process","number","the","tf_tensor","should","to","go",".","input_i",":","The","input","ordinal","of","the","tf_tensor","."]
code: def get_slices(self):
      """Yields sliced tensors and which remote pnums they should go to.
  
      Yields:
        tf_tensor: The sliced tensor.
        pnum: Which process number the tf_tensor should to go.
        input_i: The input ordinal of the tf_tensor.
      """
      with ops.device(
          _host_id_to_tf_device(self._host_id, self._external_worker)
      ):
          all_input_tensors = self._ds_iterator.get_next()
          if isinstance(all_input_tensors, tf.Tensor):
              all_input_tensors = [all_input_tensors]
          assert len(all_input_tensors) == len(
              self._all_sub_batch_pnums
          )
  
          for input_i in range(len(all_input_tensors)):
              input_tensor = all_input_tensors[input_i]
              sub_batch_pnums = self._all_sub_batch_pnums[input_i]
              mtf_input_shape = self._mtf_input_shapes[input_i]
  
              # Initialize the cache for each input_i
              self._init_slice_cache()
  
              for pnum in sub_batch_pnums:
                  # TODO(lehou): tf.slice is kinda slow. Use tf.split instead.
                  input_slice = self._slice_tensor(
                      input_tensor, mtf_input_shape, pnum
                  )
                  yield input_slice, pnum, input_i
  
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 3 for calling the SubBatchSlicer obj's underlying member method code with example usage: obj.get_slices() without return types