class_docstr: Mesh implementation for TPU using SIMD and MPI operations.
class_docstr_tok: ["Mesh","implementation","for","TPU","using","SIMD","and","MPI","operations","."]
docstr: Collective receive in groups.
  
  Each group contains the processors that differ only in mesh_axis.
  
  ```python
  group_size = self.shape[mesh_axis].size
  ```
  
  Args:
    x: a LaidOutTensor
    mesh_axis: an integer
    source_pcoord: a list of optional integers. Each element is either None
      or an integer in [0, group_size). If source_pcoord[k] is None, then the
      output for the k-th processor in each group is a zero tensor. If
      source_pcoord[k] is not None, then the output for the k-th processor in
      each group is equal to the input for the source_pcoord[k]-th processor
      in that group.
  
  Returns:
    a LaidOutTensor
code: def receive(self, x, mesh_axis, source_pcoord):
      """Collective receive in groups.
  
      Each group contains the processors that differ only in mesh_axis.
  
      ```python
      group_size = self.shape[mesh_axis].size
      ```
  
      Args:
        x: a LaidOutTensor
        mesh_axis: an integer
        source_pcoord: a list of optional integers. Each element is either None
          or an integer in [0, group_size). If source_pcoord[k] is None, then the
          output for the k-th processor in each group is a zero tensor. If
          source_pcoord[k] is not None, then the output for the k-th processor in
          each group is equal to the input for the source_pcoord[k]-th processor
          in that group.
  
      Returns:
        a LaidOutTensor
      """
      x = x.to_laid_out_tensor()
      t = x.one_slice
      source_target_pairs = []
  
      for pnum in xrange(self.size):
          coord = mtf.pnum_to_processor_coordinates(self.shape, pnum)
          k = coord[mesh_axis]
          if source_pcoord[k] is not None:
              coord[mesh_axis] = source_pcoord[k]
              source_pnum = mtf.processor_coordinates_to_pnum(
                  self.shape, coord
              )
              source_target_pairs.append(
                  [self.l2p(source_pnum), self.l2p(pnum)]
              )
  
      if not source_target_pairs:
          ret = tf.zeros_like(t, t.dtype)
      elif t.dtype in [tf.float32, tf.bfloat16, tf.int32]:
          ret = tpu_ops.collective_permute(t, source_target_pairs)
      else:
          # If t is not one of the allowed types, cast and cast back.
          ret = tf.cast(
              tpu_ops.collective_permute(
                  tf.cast(t, tf.float32), source_target_pairs
              ),
              t.dtype,
          )
  
      return self.LaidOutTensor([ret])
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import math
  import os
  import gin
  from mesh_tensorflow import ops_with_redefined_builtins
  from mesh_tensorflow import tpu_variables
  from mesh_tensorflow import utils
  from six.moves import xrange
  import tensorflow.compat.v1
  from tensorflow.python.tpu.ops import tpu_ops
label_desc: Use image node_id 10 for calling the SimdMeshImpl obj's underlying member method code with example usage: obj.receive(x, mesh_axis, source_pcoord) and returns: self