A batch of size k with k per-example gradients is numerically equivalent to iterating over k examples using batch size 1 and saving k gradients you get.
This mini-unit test in the docs of autograd-hacks should additionally clarify the relationship between per-example gradients and the “gradient”
# param.grad: gradient averaged over the batch
# param.grad1[i]: gradient with respect to example i

for param in model.parameters():
  assert(torch.allclose(param.grad1.mean(dim=0), param.grad))

Note that KFAC implementations typically inject isotropic random vectors at the output layer (ie, call output.backward(random_vec)) and compute factored approximation of resulting covariance matrix of gradients. This (in expectation) is equal to Gauss-Newton matrix, with softmax loss replaced by least-squares loss, rather than Fisher matrix.
Problem with Fisher matrix (and Gauss-Newton with softmax) is that once you start converging, those two matrices go to zero, so preconditioning by such matrix is a division by zero.
I have an example of computing KFAC factors in a slightly cleaned-up autograd-lib here:
  
      

      GitHub
  

  
    

GitHub - cybertronai/autograd-lib 4

  Contribute to cybertronai/autograd-lib development by creating an account on GitHub.