Okay, so we can flatten a matrix to make it a vector and then compute the Jacobian of that but I’m not very sure about how this big Jacobian will fit in this equation:
grad_h = grad_h_relu.dot(Jacobian(h_relu))

Here, grad_h and grad_h_relu are of size [64, 100]. Since h_relu is of size [64, 100] and relu is a pointwise operation, after flattening we will get a vector of size 6400 whose Jacobian would be a square matrix of size [6400, 6400].
This representation of combined Jacobians of all the samples can not be multiplied with grad_h_relu directly. So, I am not sure how the two matrices on the LHS would be multiplied.
One thing I can think of is, flatten the grad_h_relu to make it a vector of size 6400 then multiply with the Jacobian which is of size [6400, 6400] so we will get a vector of size 6400 which we will unflatten to make it of size [64, 100]. Is this the way we do it?
I don’t know if this forum is the right to place to discuss this or not as this discussion is not relevant to PyTorch specifically, and I apologize for bringing it up here if this is not. The doubt came while going through one of the PyTorch tutorials mentioned in the site and I thought many people using this forum would have gone through the tutorial, so thought of asking it here