class_docstr: Implements Shampoo algorithm.
  
  See details in:
  - https://arxiv.org/pdf/1802.09568.pdf
  - https://arxiv.org/pdf/2002.09018.pdf
  
  If root_inv_dist = True, assigns each parameter's preconditioners to different GPUs in a
  round-robin fashion.
  
  Uses infinity norm to evaluate residuals and errors.
  
  Args:
      params (iterable): iterable of parameters to optimize or dicts defining
          parameter groups
      lr (float): learning rate (Default: 1e-2)
      betas (Tuple[float, float]): coefficients used for computing running averages
          of gradient and its square (Default: (0.9, 1.0))
      epsilon (float): term added to the denominator to improve numerical stability (Default: 1e-12)
      use_bias_correction (bool): flag for using bias correction (Default: True)
      adam_w_mode (bool): Flag for using AdamW-style weight decay (Default: True)
      weight_decay (float): weight decay (L2 penalty) (Default: 0)
      update_freq (int): frequency for updating inverse preconditioner (Default: 100)
      init_delay (int): initial delay before starting to compute root inverse (Default: 1000)
      threshold (int): threshold for switching to diagonal preconditioner (Default: 1024)
      preconditioner_dtype (torch.dtype): data type for preconditioner (Default: torch.float)
      large_dim_method (LargeDimMethod): method for handling large scale tensors. (Default: LargeDimMethod.BLOCKING)
      root_inv_dist (bool): distributes root inverse computation across multiple GPU workers (Default: True)
      use_merge_dims (bool): merge dimensions if possible while respecting threshold. (Default: True)
      grafting_type (GraftingType): Selects grafting method. (Default: GraftingType.ADAGRAD)
      grafting_epsilon (float): Epsilon for grafting method. (Default: 1e-3)
      grafting_beta2 (float): Exponential moving average factor for grafting method. (Default: 1.0)
docstr: Broadcasts inverse preconditioners.
docstr_tok: ["Broadcasts","inverse","preconditioners","."]
code: def _broadcast_inv_preconditioners(self):
      """Broadcasts inverse preconditioners."""
  
      for group in self.param_groups:
          # get world size
          world_size = dist.get_world_size()
  
          for idx, p in enumerate(group["params"]):
              if p.grad is None:
                  continue
  
              # get state, dimensions, and order
              state = self.state[p]
              src_rank = idx % world_size
              state["preconditioners"].broadcast(src_rank)
  
file_path: bert-ultrafast/training/cramming/backend/optimizers/shampoo/shampoo.py
imports: import logging
  from typing import Tuple
  import torch
  import torch.distributed
  from torch.optim.optimizer import Optimizer
  from .shampoo_utils import BlockShampooPreconditioner, AdagradPreconditioner, ShampooPreconditioner, LargeDimMethod, GraftingType
label_desc: Use image node_id 4 for calling the Shampoo obj's underlying member method code with example usage: obj._broadcast_inv_preconditioners() without return types