Hi, I am trying to understand how to calculate the higher-order gradients using vmap and grad because it seems helpful in controlling the dimensions. However, vmap is showing error in my code:
Thanks in advance!
