Block diagonal Hessian-vector product

[exenGT]

Hi all,

I recently found that JAX can efficiently calculate the Hessian-vector product (hvp) of a scalar function f(x) and a given vector z. Suppose now my function takes a variable number of arguments: f(*x). How to calculate the hvp of f(*x) with respect to the i-th argument x_i and the corresponding given vector z_i?

[jakevdp]

I don't know of any easy/general way to do this. There is some relevant discussion here: #3801

[exenGT]

That's quite interesting! The motivation of my problem is that, for a dense Hessian H, solving Hx = -g is not easy; however, solving the preconditioned problem H_tilde x = -g where H_tilde is in block-diagonal form should be easier. It would be great if there is an efficient way to calculate this.

[exenGT]

From your answer, it seems that JAX cannot selectively apply hvp to a specific variable. Is it the case?

[exenGT]

Upon some investigation, I found that grad() has an argument argnums which specifies the index of the argument to differentiate with respect to. On the other hand, there is no argnums argument for jvp(). Is there a reason why it's not implemented?