code: def _get_module_fsdp_state_if_fully_sharded_module(
      module: nn.Module,
  ) -> Optional[_FSDPState]:
      state = _get_module_fsdp_state(module)
      if state is None:
          return None
      if state == module:  # FullyShardedDataParallel module case.
          return state
      if (
          module in state._fully_sharded_module_to_handle
      ):  # fully_shard case.
          return state
      return None
  
imports: import logging
  import traceback
  import warnings
  import weakref
  from enum import auto, Enum
  from functools import partial
  from typing import Any, Callable, cast, Dict, Generator, Iterable, List, no_type_check, Optional, Set, Tuple, Type, TYPE_CHECKING
  import torch
  import torch.distributed
  import torch.distributed.fsdp._flat_param
  import torch.nn
  from torch.distributed._composable_state import _get_module_state, _State
  from torch.distributed._tensor.device_mesh import DeviceMesh
  from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import _CHECKPOINT_PREFIX
  from torch.distributed.fsdp._fsdp_extensions import FSDPExtensions
  from torch.distributed.utils import _apply_to_tensors
  from torch.utils._mode_utils import no_dispatch
  from .api import FullOptimStateDictConfig, FullStateDictConfig, OptimStateDictConfig, ShardingStrategy, StateDictConfig, StateDictType
label_desc: Use image node_id 2 for calling a global function with example usage: _get_module_fsdp_state_if_fully_sharded_module(module) and returns: None, None, state, state