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 function returns the product of inverse of the fisher matrix and the weights gradient.
  
  Parameters:
  -----------
  m: torch.nn.Module
      Specifies the layer for which the calculation must be done on.
  p_grad_mat: torch.Tensor
      the gradients in matrix form isinstance(m.weight.grad.detach(), torch.Tensor) and i
  damping: float
      the damping factor for the calculation
  
  Return:
  -------
  torch.Tensor
      the product of inverse of the fisher matrix and the weights gradient.
docstr_tok: ["This","function","returns","the","product","of","inverse","of","the","fisher","matrix","and","the","weights","gradient",".","Parameters",":","--","--","--","--","--","-","m",":","torch.nn.Module","Specifies","the","layer","for","which","the","calculation","must","be","done","on",".","p_grad_mat",":","torch.Tensor","the","gradients","in","matrix","form","isinstance","(","m.weight.grad.detach","(",")",",","torch.Tensor",")","and","i","damping",":","float","the","damping","factor","for","the","calculation","Return",":","--","--","--","-","torch.Tensor","the","product","of","inverse","of","the","fisher","matrix","and","the","weights","gradient","."]
code: def _get_natural_grad(
      self, m: torch.nn.Module, p_grad_mat: torch.Tensor, damping: float
  ) -> List[torch.Tensor]:
      """
      This function returns the product of inverse of the fisher matrix and the weights gradient.
  
      Parameters:
      -----------
      m: torch.nn.Module
          Specifies the layer for which the calculation must be done on.
      p_grad_mat: torch.Tensor
          the gradients in matrix form isinstance(m.weight.grad.detach(), torch.Tensor) and i
      damping: float
          the damping factor for the calculation
  
      Return:
      -------
      torch.Tensor
          the product of inverse of the fisher matrix and the weights gradient.
      """
      # p_grad_mat is of output_dim * input_dim
      # inv((ss')) p_grad_mat inv(aa') = [ Q_g (1/R_g) Q_g^T ] @ p_grad_mat @ [Q_a (1/R_a) Q_a^T]
      v1 = self.Q_g[m].t() @ p_grad_mat @ self.Q_a[m]
      v2 = v1 / (
          self.d_g[m].unsqueeze(1) * self.d_a[m].unsqueeze(0) + damping
      )
      a = self.Q_g[m] @ v2 @ self.Q_a[m].t()
      if m.bias is not None:
          # we always put gradient w.r.t weight in [0]
          # and w.r.t bias in [1]
          if isinstance(
              m.weight.grad.detach(), torch.Tensor
          ) and isinstance(m.bias.grad.detach(), torch.Tensor):
              v = [a[:, :-1], a[:, -1:]]
              v[0] = v[0].view(m.weight.grad.detach().size())
              v[1] = v[1].view(m.bias.grad.detach().size())
          else:
              raise TypeError(
                  "weight.grad.detach() and bias.grad.detach() should be a Tensor"
              )
      else:
          v = [a.view(m.weight.grad.detach().size())]
  
      return v
  
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 11 for calling the KFACOptimizer obj's underlying member method code with example usage: obj._get_natural_grad(m, p_grad_mat, damping) and returns: v