Thanks!
But the Jacobian computation in autograd-lib includes the loop along the output dimension.

  
      github.com
  
  
    cybertronai/autograd-lib/blob/master/autograd_lib/autograd_lib.py#L122-L123 1

for idx in range(o):
    output.backward(torch.stack([id_mat[idx]] * n), retain_graph=(retain_graph or idx < o - 1))



  
  
    
    
  
  


bs: 128
n_inputs: 1000
n_outputs: 1000
loop: 1
-------------
autograd-lib: 2.100s
reverse mode: 1.193s
manual mode: 0.238s