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
- calculate the whole Jacobian matrix, then take out the diagonal, which takes a lot of memory.
- 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.