code: def mesh_broadcast(
      tensor: torch.Tensor,
      mesh: DeviceMesh,
      mesh_dim: int = 0,
      async_op: bool = False,
  ) -> Optional[Work]:
      """
      broadcast the tensor to a device mesh dimension. We by default
      use the first rank of the mesh dimension as the source of truth, i.e
      for a 2d mesh [[0, 1], [2, 3]], if we broadcast on mesh_dim = 1, we will
      broadcast the tensor on rank 0 to rank 0/1, and tensor on rank 2
      to rank 2/3.
  
      Args:
          tensor (torch.Tensor): tensor to broadcast.
          mesh_dim (int, optional): indicate which mesh dimension we want
              to scatter on, we by default choose the first rank on the
              mesh dimension as source of truth.
  
      Returns:
          A :class:`Work` object
      """
      # TODO: Ideally we should use the meta tensor way
      # (to register a meta kernel for the collective op)
      # so that it would avoid the communication. Need to
      # remove the check below once that is done.
      if tensor.is_meta:
          return None
      dim_group = mesh.get_group(mesh_dim)
      assert isinstance(dim_group, ProcessGroup)
      # src need to be global rank
      src_for_dim = 0
      if dim_group is not GroupMember.WORLD:
          src_for_dim = get_global_rank(dim_group, 0)
  
      return broadcast(
          tensor, src=src_for_dim, group=dim_group, async_op=async_op
      )
  
file_path: pytorch/torch/distributed/_tensor/_collective_utils.py
imports: import logging
  import math
  from typing import List, Optional
  import torch
  import torch.distributed._tensor.placement_types
  from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
  from torch.distributed.distributed_c10d import all_to_all, broadcast, get_global_rank, get_rank, get_world_size, GroupMember, ProcessGroup, scatter, Work
label_desc: Use image node_id 2 for calling a global function with example usage: mesh_broadcast(tensor, mesh, mesh_dim, async_op) and returns: broadcast, None