class_docstr: "
  This class implement the second order optimizer - KFAC, which uses Kronecker factor products of inputs and the gradients to
  get the approximate inverse fisher matrix, which is used to update the model parameters. Presently this optimizer works only
  on liner and 2D convolution layers. If you want to know more details about KFAC, please check the paper [1]_ and [2]_.
  
  References:
  -----------
  [1] Martens, James, and Roger Grosse. Optimizing Neural Networks with Kronecker-Factored Approximate Curvature.
  arXiv:1503.05671, arXiv, 7 June 2020. arXiv.org, http://arxiv.org/abs/1503.05671.
  [2] Grosse, Roger, and James Martens. A Kronecker-Factored Approximate Fisher Matrix for Convolution Layers.
  arXiv:1602.01407, arXiv, 23 May 2016. arXiv.org, http://arxiv.org/abs/1602.01407.
class_docstr_tok: ["``","This","class","implement","the","second","order","optimizer","-","KFAC",",","which","uses","Kronecker","factor","products","of","inputs","and","the","gradients","to","get","the","approximate","inverse","fisher","matrix",",","which","is","used","to","update","the","model","parameters",".","Presently","this","optimizer","works","only","on","liner","and","2D","convolution","layers",".","If","you","want","to","know","more","details","about","KFAC",",","please","check","the","paper","[","1","]","_","and","[","2","]","_.","References",":","--","--","--","--","--","-","[","1","]","Martens",",","James",",","and","Roger","Grosse",".","Optimizing","Neural","Networks","with","Kronecker-Factored","Approximate","Curvature",".","arXiv:1503.05671",",","arXiv",",","7","June","2020.","arXiv.org",",","http",":","\/\/arxiv.org\/abs\/1503.05671",".","[","2","]","Grosse",",","Roger",",","and","James","Martens",".","A","Kronecker-Factored","Approximate","Fisher","Matrix","for","Convolution","Layers",".","arXiv:1602.01407",",","arXiv",",","23","May","2016.","arXiv.org",",","http",":","\/\/arxiv.org\/abs\/1602.01407","."]
docstr: Performs clipping on the updates matrix, if the value is large. Then final value is updated in the backwards gradient data
  
  Parameters:
  -----------
  updates: Dict[torch.nn.Module,List[torch.Tensor]]
      A dicitonary containing the product of gradient and fisher inverse of each layer.
  lr: float
      learning rate of the optimizer
docstr_tok: ["Performs","clipping","on","the","updates","matrix",",","if","the","value","is","large",".","Then","final","value","is","updated","in","the","backwards","gradient","data","Parameters",":","--","--","--","--","--","-","updates",":","Dict","[","torch.nn.Module",",","List","[","torch.Tensor","]","]","A","dicitonary","containing","the","product","of","gradient","and","fisher","inverse","of","each","layer",".","lr",":","float","learning","rate","of","the","optimizer"]
code: def _kl_clip_and_update_grad(
      self,
      updates: Dict[torch.nn.Module, List[torch.Tensor]],
      lr: float,
  ):
      """
      Performs clipping on the updates matrix, if the value is large. Then final value is updated in the backwards gradient data
  
      Parameters:
      -----------
      updates: Dict[torch.nn.Module,List[torch.Tensor]]
          A dicitonary containing the product of gradient and fisher inverse of each layer.
      lr: float
          learning rate of the optimizer
      """
      # do kl clip
      vg_sum = 0.0
      for m in self.modules:
          v = updates[m]
          vg_sum += (
              (v[0] * m.weight.grad.detach() * lr**2).sum().item()
          )
          if m.bias is not None:
              vg_sum += (
                  (v[1] * m.bias.grad.detach() * lr**2).sum().item()
              )
      nu = min(1.0, math.sqrt(self.kl_clip / vg_sum))
  
      for m in self.modules:
          v = updates[m]
          if isinstance(m.weight.grad.detach(), torch.Tensor):
              m.weight.grad.detach().copy_(v[0])
              m.weight.grad.detach().mul_(nu)
          else:
              raise TypeError("weight.grad.detach() should be a Tensor")
          if m.bias is not None:
              if isinstance(m.bias.grad.detach(), torch.Tensor):
                  m.bias.grad.detach().copy_(v[1])
                  m.bias.grad.detach().mul_(nu)
              else:
                  raise TypeError(
                      "bias.grad.detach() should be a Tensor"
                  )
  
file_path: deepchem/deepchem/models/torch_models/kfac_optimizer.py
imports: import math
  from typing import Optional, Callable, Dict, Tuple, List, Union
label_desc: Use image node_id 12 for calling the KFACOptimizer obj's underlying member method code with example usage: obj._kl_clip_and_update_grad(updates, lr) without return types