code: def _init_process_group_state_for_hybrid_shard(
      state: _FSDPState,
      process_group: ProcessGroupType,
      device_mesh: DeviceMesh,
  ) -> _FSDPState:
      if device_mesh:
          if _is_valid_hybrid_shard_device_mesh(device_mesh):
              state._device_mesh = device_mesh
              # We currently only allow _inter_node_pg to be the outermost dimension, and the
              # process_group(intra_node) to be the innermost dimension.
              state._inter_node_pg = device_mesh.get_group(mesh_dim=0)
              state.process_group = device_mesh.get_group(mesh_dim=1)
          else:
              raise ValueError(
                  "Expected device_mesh to have ndim=2 "
                  f"but got {len(device_mesh.get_group())}"
              )
      elif process_group is None:
          default_group = _get_default_group()
          (
              intra_node_group,
              inter_node_group,
          ) = _init_intra_and_inter_node_groups(
              default_group, state._device_handle.device_count()
          )
          # we shard across intra-node
          state.process_group = intra_node_group
          # save _inter_node_pg to allreduce across.
          state._inter_node_pg = inter_node_group
      else:
          # Check type and assign state.process_group and state._inter_node_pg.
          if _is_valid_hybrid_shard_pg_type(process_group):
              # Assuming that user passed in as intra node group and inter node group
              # as documented.
              state.process_group, state._inter_node_pg = process_group
          else:
              raise ValueError(
                  "Expected process_group to be passed in as either None or "
                  f"Tuple[dist.ProcessGroup, dist.ProcessGroup] but got {type(process_group)}"
              )
      # Create state for allreduce
      state._inter_node_state = _get_default_comm_hook_state(
          process_group=state._inter_node_pg,
      )
      return state
  
imports: import collections
  import itertools
  import os
  import warnings
  from typing import Any, Callable, Deque, Dict, Generator, Iterable, Iterator, List, no_type_check, Optional, Set, Tuple, Union
  import torch
  import torch.distributed
  import torch.distributed.fsdp._exec_order_utils
  import torch.distributed.fsdp._traversal_utils
  import torch.distributed.fsdp.fully_sharded_data_parallel
  import torch.nn
  from torch.distributed.algorithms._comm_hooks import default_hooks
  from torch.distributed.device_mesh import _mesh_resources, DeviceMesh
  from torch.distributed.distributed_c10d import _get_default_group
  from torch.distributed.fsdp._common_utils import _FSDPDeviceHandle, _FSDPState, _get_module_fsdp_state, _is_fsdp_flattened, _named_parameters_with_duplicates, clean_tensor_name, TrainingState
  from torch.distributed.fsdp._flat_param import _FSDP_USE_FULL_PREC_IN_EVAL, FlatParameter, FlatParamHandle, HandleShardingStrategy
  from torch.distributed.fsdp._limiter_utils import _FreeEventQueue
  from torch.distributed.fsdp.api import BackwardPrefetch, CPUOffload, FullOptimStateDictConfig, FullStateDictConfig, MixedPrecision, ShardingStrategy, StateDictConfig, StateDictType
  from torch.distributed.fsdp.wrap import _Policy
  from torch.distributed.tensor.parallel.fsdp import DTensorExtensions
  from torch.distributed.utils import _sync_params_and_buffers
  from torch.utils._python_dispatch import is_traceable_wrapper_subclass
  from torch.utils.hooks import RemovableHandle
label_desc: Use image node_id 2 for calling a global function with example usage: _init_process_group_state_for_hybrid_shard(state, process_group, device_mesh) and returns: state