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.