class_docstr: This class defines the generic join context manager, which allows custom hooks to be called after a process joins.
  
  These hooks should shadow the
  collective communications of non-joined processes to prevent hanging and
  erroring and to ensure algorithmic correctness. Refer to :class:`JoinHook`
  for details about the hook definition.
  
  .. warning::
      The context manager requires each participating :class:`Joinable` to
      call the method :meth:`notify_join_context()` before its own per-
      iteration collective communications to ensure correctness.
  
  .. warning::
      The context manager requires that all ``process_group`` attributes in
      the :class:`JoinHook` objects are the same. If there are multiple
      :class:`JoinHook` objects, then the ``device`` of the first is used.
      The process group and device information is used for checking for non-
      joined processes and for notifying processes to throw an exception if
      ``throw_on_early_termination`` is enabled, both of which using an all-
      reduce.
  
  Arguments:
      joinables (List[Joinable]): a list of the participating
          :class:`Joinable` s; their hooks are iterated over in the given
          order.
  
      enable (bool): a flag enabling uneven input detection; setting to
          ``False`` disables the context manager's functionality and should
          only be set when the user knows the inputs will not be uneven
          (default: ``True``).
  
      throw_on_early_termination (bool): a flag controlling whether to throw an
          exception upon detecting uneven inputs (default: ``False``).
  
  Example::
  
      >>> import os
      >>> import torch
      >>> import torch.distributed as dist
      >>> import torch.multiprocessing as mp
      >>> # xdoctest: +SKIP
      >>> import torch.nn.parallel.DistributedDataParallel as DDP
      >>> import torch.distributed.optim.ZeroRedundancyOptimizer as ZeRO
      >>> from torch.distributed.algorithms.join import Join
      >>>
      >>> # On each spawned worker
      >>> def worker(rank):
      >>>     dist.init_process_group("nccl", rank=rank, world_size=2)
      >>>     model = DDP(torch.nn.Linear(1, 1).to(rank), device_ids=[rank])
      >>>     optim = ZeRO(model.parameters(), torch.optim.Adam, lr=0.01)
      >>>     # Rank 1 gets one more input than rank 0
      >>>     inputs = [torch.tensor([1.]).to(rank) for _ in range(10 + rank)]
      >>>     with Join([model, optim]):
      >>>         for input in inputs:
      >>>             loss = model(input).sum()
      >>>             loss.backward()
      >>>             optim.step()
      >>>     # All ranks reach here without hanging/erroring
class_docstr_tok: ["This","class","defines","the","generic","join","context","manager",",","which","allows","custom","hooks","to","be","called","after","a","process","joins",".","These","hooks","should","shadow","the","collective","communications","of","non-joined","processes","to","prevent","hanging","and","erroring","and","to","ensure","algorithmic","correctness",".","Refer","to",":","class",":","`","JoinHook","`","for","details","about","the","hook","definition",".","..","warning",":",":","The","context","manager","requires","each","participating",":","class",":","`","Joinable","`","to","call","the","method",":","meth",":","`","notify_join_context","(",")","`","before","its","own","per-","iteration","collective","communications","to","ensure","correctness",".","..","warning",":",":","The","context","manager","requires","that","all","``","process_group","``","attributes","in","the",":","class",":","`","JoinHook","`","objects","are","the","same",".","If","there","are","multiple",":","class",":","`","JoinHook","`","objects",",","then","the","``","device","``","of","the","first","is","used",".","The","process","group","and","device","information","is","used","for","checking","for","non-","joined","processes","and","for","notifying","processes","to","throw","an","exception","if","``","throw_on_early_termination","``","is","enabled",",","both","of","which","using","an","all-","reduce",".","Arguments",":","joinables","(","List","[","Joinable","]",")",":","a","list","of","the","participating",":","class",":","`","Joinable","`","s",";","their","hooks","are","iterated","over","in","the","given","order",".","enable","(","bool",")",":","a","flag","enabling","uneven","input","detection",";","setting","to","``","False","``","disables","the","context","manager","'s","functionality","and","should","only","be","set","when","the","user","knows","the","inputs","will","not","be","uneven","(","default",":","``","True","``",")",".","throw_on_early_termination","(","bool",")",":","a","flag","controlling","whether","to","throw","an","exception","upon","detecting","uneven","inputs","(","default",":","``","False","``",")",".","Example",":",":",">",">",">","import","os",">",">",">","import","torch",">",">",">","import","torch.distributed","as","dist",">",">",">","import","torch.multiprocessing","as","mp",">",">",">","#","xdoctest",":","+SKIP",">",">",">","import","torch.nn.parallel.DistributedDataParallel","as","DDP",">",">",">","import","torch.distributed.optim.ZeroRedundancyOptimizer","as","ZeRO",">",">",">","from","torch.distributed.algorithms.join","import","Join",">",">",">",">",">",">","#","On","each","spawned","worker",">",">",">","def","worker","(","rank",")",":",">",">",">","dist.init_process_group","(","``","nccl","''",",","rank=rank",",","world_size=2",")",">",">",">","model","=","DDP","(","torch.nn.Linear","(","1",",","1",")",".to","(","rank",")",",","device_ids=","[","rank","]",")",">",">",">","optim","=","ZeRO","(","model.parameters","(",")",",","torch.optim.Adam",",","lr=0.01",")",">",">",">","#","Rank","1","gets","one","more","input","than","rank","0",">",">",">","inputs","=","[","torch.tensor","(","[","1",".","]",")",".to","(","rank",")","for","_","in","range","(","10","+","rank",")","]",">",">",">","with","Join","(","[","model",",","optim","]",")",":",">",">",">","for","input","in","inputs",":",">",">",">","loss","=","model","(","input",")",".sum","(",")",">",">",">","loss.backward","(",")",">",">",">","optim.step","(",")",">",">",">","#","All","ranks","reach","here","without","hanging\/erroring"]
docstr: Extract the process group and device information from the joinables.
  
  If there are multiple joinables, then the context manager uses the
  first specified device.
  
  Preconditions:
      ``self._joinables`` is not ``None`` and is non-empty.
  
  Raises:
      ValueError
          If there are multiple conflicting ``process_group`` attributes
          among the ``Joinable`` objects.
docstr_tok: ["Extract","the","process","group","and","device","information","from","the","joinables",".","If","there","are","multiple","joinables",",","then","the","context","manager","uses","the","first","specified","device",".","Preconditions",":","``","self._joinables","``","is","not","``","None","``","and","is","non-empty",".","Raises",":","ValueError","If","there","are","multiple","conflicting","``","process_group","``","attributes","among","the","``","Joinable","``","objects","."]
code: def _extract_dist_info(self) -> None:
      r"""
      Extract the process group and device information from the joinables.
  
      If there are multiple joinables, then the context manager uses the
      first specified device.
  
      Preconditions:
          ``self._joinables`` is not ``None`` and is non-empty.
  
      Raises:
          ValueError
              If there are multiple conflicting ``process_group`` attributes
              among the ``Joinable`` objects.
      """
      process_group = None
      device = None
      for joinable in self._joinables:
          if process_group is None:
              process_group = joinable.join_process_group
          elif process_group != joinable.join_process_group:
              raise ValueError(
                  "Using join context manager with multiple process groups"
              )
          if device is None:
              device = joinable.join_device
      self._process_group = process_group
      self._rank = dist.get_rank(self._process_group)
      self._device = device
  
imports: import warnings
  from abc import ABC, abstractmethod
  from types import TracebackType
  from typing import Any, List, NamedTuple, Optional, Type
  import torch
  import torch.distributed
label_desc: Use image node_id 3 for calling the Join obj's underlying member method code with example usage: obj._extract_dist_info() without return types