class_docstr: Implements the AdaScale_ algorithm for scaling the learning rate for
  distributed and large batch size training. Can be used in combination with
  ``torch.nn.parallel.DistributedDataParallel`` and ``torch.optim.SGD``.
  
  .. _AdaScale: https://proceedings.icml.cc/static/paper_files/icml/2020/4682-Supplemental.pdf
  
  This class subclasses `Optimizer` so that `torch.optim.lr_scheduler` can
  work with it. In other words, AdaScale is intended to be a complete wrapper of an
  torch Optimizer.
  
  Note that, AdaScale does *not* help increase per-GPU batch size.
  
  There are several ways to integrate AdaScale with your training loop.
  We show two examples below.
  
  Example 1: using PyTorch's `lr_scheduler` classes.
  
  .. code-block:: python
  
      optim = AdaScale(SGD(model.parameters(), lr=0.001))
      model = DistributedDataParallel(model)
      scheduler = LambdaLR(optim, lr_lambda=...)
  
      last_epoch = 0
      done = False
      step = 0
      while not done:
          for batch in dataset:
              optim.zero_grad()
              logits = model()
              loss = criterion(logits, ...)
              loss.backward()
              step += optim.gain()
              optim.step()
              epoch = step // len(dataset)
              if epoch > last_epoch:
                  scheduler.step()
                  last_epoch = epoch
              if epoch >= MAX_EPOCHS:
                  done = True
  
  Example 2: using a custom `update_lr()` function that update the learning
  rate based on the current step count per epoch.
  
  .. code-block:: python
  
      optim = AdaScale(SGD(model.parameters(), lr=0.001))
      model = DistributedDataParallel(model)
  
      step = 0
      while step < max_steps:
          for batch in ...:
              optim.zero_grad()
              logits = model()
              loss = criterion()
              loss.backward()
              step += optim.gain()
              optim.step()
              update_lr(step)
  
  Args:
      optimizer (torch.optim.Optimizer):
          Optimizer to apply AdaScale to.
      world_size (int):
          Number of world_size for distributed training.
          If None, defaults to ``dist.get_world_size()``.
      scale (float):
          Scaling factor of the batch size from scale equals 1, e.g. using a 10x
          larger batch size (summed across all ranks with gradient accumulation)
          means a scale of 10.
          If None, defaults to ``world_size * num_gradients_to_accumulate``.
      smoothing (float):
          Smoothing factor for moving average.
          If None, it defaults to ``max(1 - (world_size * num_gradients_to_accumulate)/1000, 0)``.
          Note, for very high scale training, higher smoothing value might be needed,
          esp at the begining of the training. Therefore, if your scale is close to or larger
          than 1000, try experimenting with smoothing value > 0 if the final accuracy is poor.
      num_gradients_to_accumulate (int):
          Number of passes that we accumulate gradients locally
          between each optimizer step. This can be changed during
          training as long as the train loop changes gradient accumulation
          accordingly.
          The loss in each pass can be either scaled or unscaled. See `is_scaled_loss` below.
          Default to 1, which does not accumulate gradients.
      is_scaled_loss (bool):
          If True, assume that the loss is scaled by `num_gradients_to_accumulate`.
          If False, the loss is not scaled.
          Default: True.
      debias_ewma (bool):
          (experimental) Use debias exponential moving average
          for smoothing and mu and sigma variables. False will
          use the method in the paper's Appendix B.3.
          Default: True, which is what have been validated so far.
docstr: Helper function for computing average intra correlation among gradients on different GPUs.
  This should be called under `model.no_sync()` context.
docstr_tok: ["Helper","function","for","computing","average","intra","correlation","among","gradients","on","different","GPUs",".","This","should","be","called","under","`","model.no_sync","(",")","`","context","."]
code: def _compute_intra_grad_corr_mean(self) -> torch.Tensor:
      """
      Helper function for computing average intra correlation among gradients on different GPUs.
      This should be called under `model.no_sync()` context.
      """
      assert self._world_size > 1, "Only for distributed training"
      flat_grad = self._gather_flat_grad()
      corr_mean = torch.tensor(0.0).cuda()
      if dist.get_rank() == 0:
          size = flat_grad.numel()
          gathered_tensors = [
              torch.zeros(size, device=0)
              for _ in range(self._world_size)
          ]
          dist.gather(flat_grad, gather_list=gathered_tensors, dst=0)
          # the following requires torch 1.10+
          corr = torch.stack(gathered_tensors).corrcoef()  # type: ignore
          # pick out the upper triangular part of the correlation matrix
          corr = corr[
              torch.triu(torch.ones_like(corr), diagonal=1) == 1
          ]
          corr_mean = corr.mean()
      else:
          dist.gather(flat_grad, gather_list=None, dst=0)
      dist.broadcast(corr_mean, src=0)
      return corr_mean
  
imports: import functools
  from typing import TYPE_CHECKING, Any, Dict, List, Optional, Type
  import numpy
  import torch
  from torch.autograd import Variable
  import torch.distributed
  from torch.optim import SGD, Optimizer
label_desc: Use image node_id 14 for calling the AdaScale obj's underlying member method code with example usage: obj._compute_intra_grad_corr_mean() and returns: corr_mean