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: This is the function that gets called in each step of the optimizer to update the weights and biases of the model.
  
  Parameters:
  -----------
  closure: Callable, optional(default: None)
      an optional customizable function to be passed which can be used to clear the gradients and other compute loss for every step.
docstr_tok: ["This","is","the","function","that","gets","called","in","each","step","of","the","optimizer","to","update","the","weights","and","biases","of","the","model",".","Parameters",":","--","--","--","--","--","-","closure",":","Callable",",","optional","(","default",":","None",")","an","optional","customizable","function","to","be","passed","which","can","be","used","to","clear","the","gradients","and","other","compute","loss","for","every","step","."]
code: def step(self, closure: Optional[Callable] = None):
      """
      This is the function that gets called in each step of the optimizer to update the weights and biases of the model.
  
      Parameters:
      -----------
      closure: Callable, optional(default: None)
          an optional customizable function to be passed which can be used to clear the gradients and other compute loss for every step.
      """
      group = self.param_groups[0]
      lr = group["lr"]
      damping = group["damping"]
      updates = {}
      for m in self.modules:
          if self.steps % self.TInv == 0:
              self._update_inv(m)
          p_grad_mat = self._get_matrix_form_grad(m)
          v = self._get_natural_grad(m, p_grad_mat, damping)
          updates[m] = v
      self._kl_clip_and_update_grad(updates, lr)
  
      self._step(closure)
      self.steps += 1
  
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 14 for calling the KFACOptimizer obj's underlying member method code with example usage: obj.step(closure) without return types