code: def _p2p_comm(
      tensor_send_next: torch.Tensor,
      recv_prev: bool,
      peer: int,
      group: ProcessGroup,
      comm_dtype: torch.dtype = torch.float16,
  ):
      """
      Send and recv tensor using P2P communication, used when pipeline size is 2 to solve the race communication.
  
      Args:
          tensor_send_next (torch.Tensor): tensor to be sent to next stage
          recv_prev (bool): whether to receive tensor from previous stage
          peer (int): rank of the peer
          group (ProcessGroup): process group
          comm_dtype (torch.dtype): dtype of the tensor to be sent
  
      Returns:
          torch.Tensor: tensor received from previous stage
      """
      # send and recv shape
      send_next_shape = None
      recv_prev_shape = None
  
      if tensor_send_next is not None:
          send_next_shape = torch.tensor(
              tensor_send_next.size(),
              device=torch.cuda.current_device(),
              dtype=torch.int64,
          )
      if recv_prev:
          recv_prev_shape = torch.empty(
              (3), device=torch.cuda.current_device(), dtype=torch.int64
          )
  
      ops = []
      if send_next_shape is not None:
          send_next_op = dist.P2POp(
              dist.isend, send_next_shape, peer=peer, group=group
          )
          ops.append(send_next_op)
      if recv_prev_shape is not None:
          recv_prev_op = dist.P2POp(
              dist.irecv,
              recv_prev_shape,
              peer=peer,
              group=group,
          )
          ops.append(recv_prev_op)
      if len(ops) > 0:
          reqs = dist.batch_isend_irecv(ops)
          for req in reqs:
              req.wait()
  
      if recv_prev_shape is not None:
          recv_prev_shape = recv_prev_shape.tolist()
  
      # send and recv data
      tensor_recv_prev = None
      if recv_prev:
          tensor_recv_prev = torch.empty(
              recv_prev_shape,
              device=torch.cuda.current_device(),
              dtype=comm_dtype,
          )
  
      ops = []
      if tensor_send_next is not None:
          send_next_op = dist.P2POp(
              dist.isend,
              tensor_send_next,
              peer=peer,
              group=group,
          )
          ops.append(send_next_op)
      if tensor_recv_prev is not None:
          recv_prev_op = dist.P2POp(
              dist.irecv,
              tensor_recv_prev,
              peer=peer,
              group=group,
          )
          ops.append(recv_prev_op)
      if len(ops) > 0:
          reqs = dist.batch_isend_irecv(ops)
          for req in reqs:
              req.wait()
      return tensor_recv_prev
  
imports: import io
  import pickle
  import re
  from collections import namedtuple
  from typing import Any, Callable, List, Optional, Tuple, Union
  import torch
  import torch.distributed
  from packaging.version import Version
  from torch.distributed import ProcessGroup
  from torch.distributed import distributed_c10d
  from torch.utils._pytree import tree_flatten, tree_unflatten
  from .stage_manager import PipelineStageManager
label_desc: Use image node_id 13 for calling a global function with example usage: _p2p_comm(tensor_send_next, recv_prev, peer, group, comm_dtype) and returns: tensor_recv_prev