code: def wrap(module: nn.Module, **wrap_overrides: Any) -> nn.Module:
      """
      Annotate that a module should be wrapped. Annotated modules will only be
      wrapped if inside of an :func:`enable_wrap` context manager. This allows
      a module to be initialized both with and without a wrapper without code
      change.
  
      Both wrapper_cls and wrapper_config can be taken from 3 sources with
      increasing priority:
  
          1. ConfigAutoWrap's context
          2. module.wrapper_config
          3. wrap_overrides argument of this function
  
      Usage::
  
          with enable_wrap(wrapper_cls=FSDP, **fsdp_config):
              # Wraps layer in FSDP by default if within context
              self.l1 = wrap(torch.nn.Linear(5, 5))
  
      Args:
          module (nn.Module): module to wrap (if in :func:`enable_wrap` context)
          **wrap_overrides: configuration overrides that will take priority over
              the values provided by the :func:`enable_wrap` context
      """
      if ConfigAutoWrap.in_autowrap_context:
          module_overrides = {}
          if hasattr(module, "wrapper_config"):
              module_overrides = module.wrapper_config
              assert isinstance(module_overrides, dict)
          wrap_overrides = {
              **ConfigAutoWrap.kwargs,
              **module_overrides,
              **wrap_overrides,
          }
          assert ConfigAutoWrap.wrapper_cls is not None
          if ConfigAutoWrap.move_module_cuda_half:
              module = module.cuda().half()
          return ConfigAutoWrap.wrapper_cls(module, **wrap_overrides)
      return module
  
file_path: fairscale/fairscale/nn/wrap/auto_wrap.py
imports: import contextlib
  from typing import Any, Callable, Dict, Generator, Optional, Set, Tuple, Type, cast
  import torch.nn
label_desc: Use image node_id 4 for calling a global function with example usage: wrap(module) and returns: module, ConfigAutoWrap