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: Compute the covariance matrix of the G matrix (the gradient of the layer).
  
  Parameters:
  -----------
  g: torch.Tensor
      It is the gradient of the layer for which the covariance matrix should be calculated.
  layer: torch.nn.Module
      It specifies the type of layer from which the output of the layer is taken.
  
  Returns:
  --------
  torch.Tensor
      The covariance matrix of the G matrix.
docstr_tok: ["Compute","the","covariance","matrix","of","the","G","matrix","(","the","gradient","of","the","layer",")",".","Parameters",":","--","--","--","--","--","-","g",":","torch.Tensor","It","is","the","gradient","of","the","layer","for","which","the","covariance","matrix","should","be","calculated",".","layer",":","torch.nn.Module","It","specifies","the","type","of","layer","from","which","the","output","of","the","layer","is","taken",".","Returns",":","--","--","--","--","torch.Tensor","The","covariance","matrix","of","the","G","matrix","."]
code: def compute_cov_g(
      self, g: torch.Tensor, layer: torch.nn.Module
  ) -> torch.Tensor:
      """
      Compute the covariance matrix of the G matrix (the gradient of the layer).
  
      Parameters:
      -----------
      g: torch.Tensor
          It is the gradient of the layer for which the covariance matrix should be calculated.
      layer: torch.nn.Module
          It specifies the type of layer from which the output of the layer is taken.
  
      Returns:
      --------
      torch.Tensor
          The covariance matrix of the G matrix.
      """
      if isinstance(layer, torch.nn.Linear):
          batch_size = g.size(0)
          if self.batch_averaged:
              cov_g = g.t() @ (g * batch_size)
          else:
              cov_g = g.t() @ (g / batch_size)
  
      elif isinstance(layer, torch.nn.Conv2d):
          spatial_size = g.size(2) * g.size(3)
          batch_size = g.shape[0]
          g = g.transpose(1, 2).transpose(2, 3)
          g = self.try_contiguous(g)
          g = g.view(-1, g.size(-1))
          if self.batch_averaged:
              g = g * batch_size
          g = g * spatial_size
          cov_g = g.t() @ (g / g.size(0))
  
      return cov_g
  
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 5 for calling the KFACOptimizer obj's underlying member method code with example usage: obj.compute_cov_g(g, layer) and returns: cov_g