class_docstr: Handles input pipeline for SimdMeshImpl.
class_docstr_tok: ["Handles","input","pipeline","for","SimdMeshImpl","."]
docstr: Returns the pnum_map according to mtf_shape.
  
  Args:
    mtf_shape: A mtf.Shape object.
  Returns:
    A numpy array pnum_map. For the i-th sub-batch, pnum_map[i] is a numpy
    array containing all pnums that tensor slices of the i-th sub-batch
    will be send to.
docstr_tok: ["Returns","the","pnum_map","according","to","mtf_shape",".","Args",":","mtf_shape",":","A","mtf.Shape","object",".","Returns",":","A","numpy","array","pnum_map",".","For","the","i-th","sub-batch",",","pnum_map","[","i","]","is","a","numpy","array","containing","all","pnums","that","tensor","slices","of","the","i-th","sub-batch","will","be","send","to","."]
code: def _get_pnum_map(self, mtf_shape):
      """Returns the pnum_map according to mtf_shape.
  
      Args:
        mtf_shape: A mtf.Shape object.
      Returns:
        A numpy array pnum_map. For the i-th sub-batch, pnum_map[i] is a numpy
        array containing all pnums that tensor slices of the i-th sub-batch
        will be send to.
      """
      s_shape = self._simd_mesh_impl.slice_shape(mtf_shape)
      shape_list = [
          dim_size // s_dim_size
          for dim_size, s_dim_size in zip(
              mtf_shape.to_integer_list, s_shape
          )
      ]
  
      pnum_map_shape = shape_list + [
          self._p_dev.num_cores // np.prod(shape_list)
      ]
      assert np.prod(pnum_map_shape) == self._p_dev.num_cores
  
      # Initialize the pnum_map to _NONE_PNUM.
      pnum_map = np.empty(pnum_map_shape, dtype=object)
      pnum_map[:] = _NONE_PNUM
  
      for pnum in range(self._p_dev.num_cores):
          s_begin = self._simd_mesh_impl.slice_begin(mtf_shape, pnum)
          coord = [
              dim_size // s_dim_size
              for dim_size, s_dim_size in zip(s_begin, s_shape)
          ]
          # put pnum in pnum_map[coord]
          pnum_array_ref = pnum_map[tuple(coord)]
          for idx, value in enumerate(pnum_array_ref):
              if value is _NONE_PNUM:
                  pnum_array_ref[idx] = pnum
                  break
  
      tf.logging.info("MTF pnum_map: {}".format(pnum_map))
      assert _NONE_PNUM not in pnum_map
      return pnum_map
  
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 5 for calling the SimdMeshImplInputReader obj's underlying member method code with example usage: obj._get_pnum_map(mtf_shape) and returns: pnum_map