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: Returns the gradient of the layer in a matrix form
  
  Parameter:
  ----------
  m: torch.nn.Module
      the layer for which the gradient must be calculated
  
  Return:
  -------
  torch.tensor
      a matrix form of the gradient. it should be a [output_dim, input_dim] matrix.
docstr_tok: ["Returns","the","gradient","of","the","layer","in","a","matrix","form","Parameter",":","--","--","--","--","--","m",":","torch.nn.Module","the","layer","for","which","the","gradient","must","be","calculated","Return",":","--","--","--","-","torch.tensor","a","matrix","form","of","the","gradient",".","it","should","be","a","[","output_dim",",","input_dim","]","matrix","."]
code: def _get_matrix_form_grad(m: torch.nn.Module):
      """
      Returns the gradient of the layer in a matrix form
  
      Parameter:
      ----------
      m: torch.nn.Module
          the layer for which the gradient must be calculated
  
      Return:
      -------
      torch.tensor
          a matrix form of the gradient. it should be a [output_dim, input_dim] matrix.
      """
      if isinstance(m, torch.nn.Conv2d):
          assert isinstance(m.weight.grad, torch.Tensor)
          p_grad_mat = m.weight.grad.detach().view(
              m.weight.grad.detach().size(0), -1
          )  # n_filters * (in_c * kw * kh)
      elif isinstance(m, torch.nn.Linear):
          assert isinstance(m.weight.grad, torch.Tensor)
          p_grad_mat = m.weight.grad.detach()
      else:
          raise NotImplementedError(
              "KFAC optimizer currently support only Linear and Conv2d layers"
          )
  
      if m.bias is not None:
          if isinstance(m.bias.grad.detach(), torch.Tensor):
              p_grad_mat = torch.cat(
                  [p_grad_mat, m.bias.grad.detach().view(-1, 1)], 1
              )
          else:
              raise TypeError("bias.grad.detach() should be a Tensor")
      return p_grad_mat
  
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 10 for calling the KFACOptimizer obj's underlying member method code with example usage: obj._get_matrix_form_grad(m) and returns: p_grad_mat