Diagonal of jacobian

autograd
1

I’d like to calculate only the diagonal of the jacobian matrix.
[d z_1/d x_1, d z_2/d x_2, …, d z_n/d x_n]
This operator is widely used in the diffusion model.
The alternatives are

  1. calculate the whole Jacobian matrix, then take out the diagonal, which takes a lot of memory.
  2. I can loop over the entries, but slow.
    So if I want to develop a specialized operator like this, please instruct how I can do so.
1 Like