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: Updates the backward gradient of the layer using exponentially weighted averages of the layer input.
  
  Parameters:
  -----------
  module: torch.nn.Module
      specifies the layer for which the gradient should be taken
  input: torch.Tensor
      the gradient matrix which should get updated
docstr_tok: ["Updates","the","backward","gradient","of","the","layer","using","exponentially","weighted","averages","of","the","layer","input",".","Parameters",":","--","--","--","--","--","-","module",":","torch.nn.Module","specifies","the","layer","for","which","the","gradient","should","be","taken","input",":","torch.Tensor","the","gradient","matrix","which","should","get","updated"]
code: def _save_grad_output(
      self,
      module: torch.nn.Module,
      grad_input: torch.Tensor,
      grad_output: torch.Tensor,
  ):
      """
      Updates the backward gradient of the layer using exponentially weighted averages of the layer input.
  
      Parameters:
      -----------
      module: torch.nn.Module
          specifies the layer for which the gradient should be taken
      input: torch.Tensor
          the gradient matrix which should get updated
      """
      # Accumulate statistics for Fisher matrices
      if self.steps % self.TCov == 0:
          gg = self.compute_cov_g(grad_output[0].data, module)
          # Initialize buffers
          if self.steps == 0:
              self.m_gg[module] = torch.diag(
                  gg.new(gg.size(0)).fill_(1)
              )
          self.m_gg[module] *= self.stat_decay + gg * (
              1 - self.stat_decay
          )
  
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 7 for calling the KFACOptimizer obj's underlying member method code with example usage: obj._save_grad_output(module, grad_input, grad_output) without return types