class_docstr: A shard aware Grad Scaler which enables loss scaling with/without cpu_offload. This is a
  slight modification of the pytorch grad scaler.
  https://pytorch.org/docs/stable/amp.html#torch.cuda.amp.GradScaler
class_docstr_tok: ["A","shard","aware","Grad","Scaler","which","enables","loss","scaling","with\/without","cpu_offload",".","This","is","a","slight","modification","of","the","pytorch","grad","scaler",".","https",":","\/\/pytorch.org\/docs\/stable\/amp.html","#","torch.cuda.amp.GradScaler"]
code: def _unscale_grads_(  # type: ignore
      self,
      optimizer: SGD,
      inv_scale: torch.Tensor,
      found_inf: torch.Tensor,
      allow_fp16: bool = True,
  ) -> Dict[torch.device, torch.Tensor]:
      per_device_inv_scale = _GeneralMultiDeviceReplicator(inv_scale)
      per_device_found_inf = _GeneralMultiDeviceReplicator(found_inf)
  
      # To set up _amp_foreach_non_finite_check_and_unscale_, split grads by device and dtype.
      # There could be hundreds of grads, so we'd like to iterate through them just once.
      # However, we don't know their devices or dtypes in advance.
  
      # https://stackoverflow.com/questions/5029934/defaultdict-of-defaultdict
      # Google says mypy struggles with defaultdicts type annotations.
      per_device_and_dtype_grads = defaultdict(lambda: defaultdict(list))  # type: ignore[var-annotated]
      with torch.no_grad():
          for group in optimizer.param_groups:
              for param in group["params"]:
                  if param.grad is None:
                      continue
                  if (
                      not allow_fp16
                  ) and param.grad.dtype == torch.float16:
                      raise ValueError(
                          "Attempting to unscale FP16 gradients."
                      )
                  if param.grad.is_sparse:
                      # is_coalesced() == False means the sparse grad has values with duplicate indices.
                      # coalesce() deduplicates indices and adds all values that have the same index.
                      # For scaled fp16 values, there's a good chance coalescing will cause overflow,
                      # so we should check the coalesced _values().
                      if param.grad.dtype is torch.float16:
                          param.grad = param.grad.coalesce()
                      to_unscale = param.grad._values()
                  else:
                      to_unscale = param.grad
  
                  # TODO: is there a way to split by device and dtype without appending in the inner loop?
                  per_device_and_dtype_grads[to_unscale.device][
                      to_unscale.dtype
                  ].append(to_unscale)
  
          for (
              device,
              per_dtype_grads,
          ) in per_device_and_dtype_grads.items():
              for grads in per_dtype_grads.values():
                  if grads[0].device.type == "cpu":
                      self._foreach_non_finite_check_and_unscale_cpu_(
                          grads,
                          per_device_found_inf.get(device),
                          per_device_inv_scale.get(device),
                      )
                  else:
                      torch._amp_foreach_non_finite_check_and_unscale_(  # type: ignore
                          grads,
                          per_device_found_inf.get(device),
                          per_device_inv_scale.get(device),
                      )
  
      return per_device_found_inf._per_device_tensors
  
imports: from collections import defaultdict
  import collections.abc
  from enum import Enum
  import logging
  from typing import Any, Dict, List, Optional, Union
  import warnings
  import torch
  from torch.cuda import FloatTensor
  from torch.cuda.amp.common import amp_definitely_not_available
  from torch.cuda.amp.grad_scaler import GradScaler
  import torch.distributed
  from torch.optim import Optimizer
  from torch.optim.sgd import SGD
  from fairscale.internal import torch_version
label_desc: Use image node_id 4 for calling the ShardedGradScaler obj's underlying member method code with example usage: obj._unscale_grads_(optimizer, inv_scale, found_inf, allow_fp16) and returns: per_device_found_inf